From f9d1e34277ef70396a54e796f751cde0118c00dd Mon Sep 17 00:00:00 2001 From: ZoeLeibowitz Date: Wed, 13 Dec 2023 10:46:24 +0000 Subject: [PATCH 01/84] compiler: Add PETSc module --- .github/workflows/pytest-petsc.yml | 75 +++ conftest.py | 16 +- devito/ir/equations/equation.py | 10 +- devito/ir/iet/algorithms.py | 4 + devito/ir/iet/efunc.py | 4 +- devito/ir/iet/nodes.py | 59 +- devito/ir/iet/visitors.py | 14 +- devito/operator/operator.py | 14 +- devito/passes/iet/definitions.py | 14 +- devito/passes/iet/engine.py | 7 +- devito/passes/iet/mpi.py | 4 +- devito/petsc/__init__.py | 1 + devito/petsc/clusters.py | 27 + devito/petsc/iet/__init__.py | 1 + devito/petsc/iet/nodes.py | 39 ++ devito/petsc/iet/passes.py | 136 ++++ devito/petsc/iet/routines.py | 1003 ++++++++++++++++++++++++++++ devito/petsc/iet/utils.py | 22 + devito/petsc/solve.py | 241 +++++++ devito/petsc/types/__init__.py | 3 + devito/petsc/types/array.py | 120 ++++ devito/petsc/types/object.py | 211 ++++++ devito/petsc/types/types.py | 131 ++++ devito/petsc/utils.py | 53 ++ devito/symbolics/extended_sympy.py | 4 + devito/symbolics/inspection.py | 9 +- devito/tools/dtypes_lowering.py | 1 + devito/types/equation.py | 4 + devito/types/object.py | 29 +- docker/Dockerfile.cpu | 4 +- docker/Dockerfile.devito | 24 +- tests/test_iet.py | 30 +- tests/test_petsc.py | 774 +++++++++++++++++++++ tests/test_symbolics.py | 18 +- 34 files changed, 3067 insertions(+), 39 deletions(-) create mode 100644 .github/workflows/pytest-petsc.yml create mode 100644 devito/petsc/__init__.py create mode 100644 devito/petsc/clusters.py 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/iet/routines.py create mode 100644 devito/petsc/iet/utils.py create mode 100644 devito/petsc/solve.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 create mode 100644 devito/petsc/utils.py create mode 100644 tests/test_petsc.py diff --git a/.github/workflows/pytest-petsc.yml b/.github/workflows/pytest-petsc.yml new file mode 100644 index 0000000000..20bfc645d1 --- /dev/null +++ b/.github/workflows/pytest-petsc.yml @@ -0,0 +1,75 @@ +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 + pull_request: + branches: + - master + - FieldFromPointer + +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 --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 }} diff --git a/conftest.py b/conftest.py index 5a1bc481b9..fc6090c604 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', - 'noadvisor', 'cpu64-arm', 'cpu64-icpx', 'chkpnt'}) + 'noadvisor', 'cpu64-arm', 'cpu64-icpx', 'chkpnt', 'petsc'}) accepted.update({'nodevice'}) unknown = sorted(set(items) - accepted) if unknown: @@ -93,6 +94,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/equations/equation.py b/devito/ir/equations/equation.py index ada1c23f22..67ed9269a4 100644 --- a/devito/ir/equations/equation.py +++ b/devito/ir/equations/equation.py @@ -9,10 +9,12 @@ 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 import (Eq, Inc, ReduceMax, ReduceMin, + relational_min) +from devito.types.equation import InjectSolveEq __all__ = ['LoweredEq', 'ClusterizedEq', 'DummyEq', 'OpInc', 'OpMin', 'OpMax', - 'identity_mapper'] + 'identity_mapper', 'OpInjectSolve'] class IREq(sympy.Eq, Pickable): @@ -102,7 +104,8 @@ def detect(cls, expr): reduction_mapper = { Inc: OpInc, ReduceMax: OpMax, - ReduceMin: OpMin + ReduceMin: OpMin, + InjectSolveEq: OpInjectSolve } try: return reduction_mapper[type(expr)] @@ -119,6 +122,7 @@ def detect(cls, expr): OpInc = Operation('+') OpMax = Operation('max') OpMin = Operation('min') +OpInjectSolve = Operation('solve') identity_mapper = { diff --git a/devito/ir/iet/algorithms.py b/devito/ir/iet/algorithms.py index 0b57b876f7..9d2db185db 100644 --- a/devito/ir/iet/algorithms.py +++ b/devito/ir/iet/algorithms.py @@ -3,6 +3,8 @@ from devito.ir.iet import (Expression, Increment, Iteration, List, Conditional, SyncSpot, Section, HaloSpot, ExpressionBundle) from devito.tools import timed_pass +from devito.petsc.types import LinearSolveExpr +from devito.petsc.iet.utils import petsc_iet_mapper __all__ = ['iet_build'] @@ -24,6 +26,8 @@ def iet_build(stree): for e in i.exprs: if e.is_Increment: exprs.append(Increment(e)) + elif isinstance(e.rhs, LinearSolveExpr): + 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) diff --git a/devito/ir/iet/efunc.py b/devito/ir/iet/efunc.py index 10aa8920e6..064f02c487 100644 --- a/devito/ir/iet/efunc.py +++ b/devito/ir/iet/efunc.py @@ -1,6 +1,6 @@ from functools import cached_property -from devito.ir.iet.nodes import Call, Callable +from devito.ir.iet.nodes import Call, Callable, FixedArgsCallable from devito.ir.iet.utils import derive_parameters from devito.symbolics import uxreplace from devito.tools import as_tuple @@ -131,7 +131,7 @@ class AsyncCall(Call): pass -class ThreadCallable(Callable): +class ThreadCallable(FixedArgsCallable): """ A Callable executed asynchronously by a thread. diff --git a/devito/ir/iet/nodes.py b/devito/ir/iet/nodes.py index 60646066b7..a066dc60ca 100644 --- a/devito/ir/iet/nodes.py +++ b/devito/ir/iet/nodes.py @@ -21,7 +21,7 @@ ctypes_to_cstr) from devito.types.basic import (AbstractFunction, AbstractSymbol, Basic, Indexed, Symbol) -from devito.types.object import AbstractObject, LocalObject +from devito.types.object import AbstractObject, LocalObject, LocalCompositeObject __all__ = ['Node', 'MultiTraversable', 'Block', 'Expression', 'Callable', 'Call', 'ExprStmt', 'Conditional', 'Iteration', 'List', 'Section', @@ -30,7 +30,7 @@ 'Increment', 'Return', 'While', 'ListMajor', 'ParallelIteration', 'ParallelBlock', 'Dereference', 'Lambda', 'SyncSpot', 'Pragma', 'DummyExpr', 'BlankLine', 'ParallelTree', 'BusyWait', 'UsingNamespace', - 'Using', 'CallableBody', 'Transfer'] + 'Using', 'CallableBody', 'Transfer', 'Callback', 'FixedArgsCallable'] # First-class IET nodes @@ -759,6 +759,15 @@ def defines(self): return self.all_parameters +class FixedArgsCallable(Callable): + + """ + A Callable class that enforces a fixed function signature. + """ + + pass + + class CallableBody(MultiTraversable): """ @@ -1037,8 +1046,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`. * `pointer` is a Symbol with its _C_ctype deriving from ct._Pointer, and `pointee` is a Symbol representing the dereferenced value. """ @@ -1061,7 +1070,8 @@ def functions(self): def expr_symbols(self): ret = [] if self.pointer.is_Symbol: - assert issubclass(self.pointer._C_ctype, ctypes._Pointer), \ + assert (isinstance(self.pointer, LocalCompositeObject) or + issubclass(self.pointer._C_ctype, ctypes._Pointer)), \ "Scalar dereference must have a pointer ctype" ret.extend([self.pointer._C_symbol, self.pointee._C_symbol]) elif self.pointer.is_PointerArray or self.pointer.is_TempFunction: @@ -1136,6 +1146,45 @@ def defines(self): return tuple(self.parameters) +class Callback(Call): + """ + Base class for special callback types. + + 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. + + 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. + """ + # TODO: Create a common base class for Call and Callback to avoid + # having arguments=None here + 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 3817c4e39d..5e3a8c0efa 100644 --- a/devito/ir/iet/visitors.py +++ b/devito/ir/iet/visitors.py @@ -265,7 +265,7 @@ def _gen_value(self, obj, mode=1, masked=()): strtype = f'{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 @@ -644,6 +644,9 @@ def visit_Lambda(self, o): top = c.Line(f"[{', '.join(captures)}]({', '.join(decls)}){''.join(extra)}") return LambdaCollection([top, c.Block(body)]) + def visit_Callback(self, o, nested_call=False): + return CallbackArg(o) + def visit_HaloSpot(self, o): body = flatten(self._visit(i) for i in o.children) return c.Collection(body) @@ -1469,3 +1472,12 @@ def sorted_efuncs(efuncs): CommCallable: 1 } return sorted_priority(efuncs, priority) + + +class CallbackArg(c.Generable): + + def __init__(self, callback): + self.callback = callback + + def generate(self): + yield self.callback.callback_form diff --git a/devito/operator/operator.py b/devito/operator/operator.py index 0a9f9db62f..4f08232050 100644 --- a/devito/operator/operator.py +++ b/devito/operator/operator.py @@ -35,7 +35,8 @@ from devito.types import (Buffer, Evaluable, host_layer, device_layer, disk_layer) from devito.types.dimension import Thickness - +from devito.petsc.iet.passes import lower_petsc +from devito.petsc.clusters import petsc_preprocess __all__ = ['Operator'] @@ -264,6 +265,9 @@ def _lower(cls, expressions, **kwargs): kwargs.setdefault('langbb', cls._Target.langbb()) kwargs.setdefault('printer', cls._Target.Printer) + # TODO: To be updated based on changes in #2509 + kwargs.setdefault('concretize_mapper', {}) + expressions = as_tuple(expressions) # Enable recursive lowering @@ -381,6 +385,9 @@ def _lower_clusters(cls, expressions, profiler=None, **kwargs): # Build a sequence of Clusters from a sequence of Eqs clusters = clusterize(expressions, **kwargs) + # 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) @@ -478,6 +485,9 @@ def _lower_iet(cls, uiet, profiler=None, **kwargs): # Lower IET to a target-specific IET graph = Graph(iet, **kwargs) + + lower_petsc(graph, **kwargs) + graph = cls._specialize_iet(graph, **kwargs) # Instrument the IET for C-level profiling @@ -512,7 +522,7 @@ 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) + ret.update(dimensions) ret = tuple(sorted(ret, key=attrgetter('name'))) diff --git a/devito/passes/iet/definitions.py b/devito/passes/iet/definitions.py index 3cc5446d57..8f65899762 100644 --- a/devito/passes/iet/definitions.py +++ b/devito/passes/iet/definitions.py @@ -19,7 +19,7 @@ SizeOf, VOID, pow_to_mul) from devito.tools import as_mapper, as_list, as_tuple, filter_sorted, flatten from devito.types import (Array, ComponentAccess, CustomDimension, DeviceMap, - DeviceRM, Eq, Symbol) + DeviceRM, Eq, Symbol, IndexedData) __all__ = ['DataManager', 'DeviceAwareDataManager', 'Storage'] @@ -333,6 +333,10 @@ def _inject_definitions(self, iet, storage): init = self.langbb['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) @@ -407,11 +411,10 @@ def place_definitions(self, iet, globs=None, **kwargs): # Track, to be handled by the EntryFunction being a global obj! globs.add(i) - elif i.is_ObjectArray: - self._alloc_object_array_on_low_lat_mem(iet, i, storage) - elif i.is_PointerArray: self._alloc_pointed_array_on_high_bw_mem(iet, i, storage) + else: + self._alloc_object_array_on_low_lat_mem(iet, i, storage) # Handle postponed global objects includes = set() @@ -447,7 +450,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)) - bases = sorted({i.base for i in indexeds}, 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/passes/iet/engine.py b/devito/passes/iet/engine.py index b9b5bf15d4..a88cf58d45 100644 --- a/devito/passes/iet/engine.py +++ b/devito/passes/iet/engine.py @@ -3,7 +3,7 @@ from devito.ir.iet import (Call, ExprStmt, Iteration, SyncSpot, AsyncCallable, FindNodes, FindSymbols, MapNodes, MetaCall, Transformer, - EntryFunction, ThreadCallable, Uxreplace, + EntryFunction, FixedArgsCallable, Uxreplace, derive_parameters) from devito.ir.support import SymbolRegistry from devito.mpi.distributed import MPINeighborhood @@ -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 @@ -602,12 +603,12 @@ def update_args(root, efuncs, dag): foo(..., z) : root(x, z) """ - if isinstance(root, ThreadCallable): + if isinstance(root, FixedArgsCallable): return efuncs # The parameters/arguments lists may have changed since a pass may have: # 1) introduced a new symbol - new_params = derive_parameters(root) + new_params = derive_parameters(root, drop_locals=True) # 2) defined a symbol for which no definition was available yet (e.g. # via a malloc, or a Dereference) diff --git a/devito/passes/iet/mpi.py b/devito/passes/iet/mpi.py index 930cc108b2..180d2231a6 100644 --- a/devito/passes/iet/mpi.py +++ b/devito/passes/iet/mpi.py @@ -72,7 +72,9 @@ def _hoist_invariant(iet): """ # Precompute scopes to save time - scopes = {i: Scope([e.expr for e in v]) for i, v in MapNodes().visit(iet).items()} + scopes = {} + 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/__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/petsc/clusters.py b/devito/petsc/clusters.py new file mode 100644 index 0000000000..0a6679b2dd --- /dev/null +++ b/devito/petsc/clusters.py @@ -0,0 +1,27 @@ +from devito.tools import timed_pass +from devito.petsc.types import LinearSolveExpr + + +@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 + distinct iteration loops. + """ + 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)) + else: + processed.append(c) + return processed diff --git a/devito/petsc/iet/__init__.py b/devito/petsc/iet/__init__.py new file mode 100644 index 0000000000..beb6d1f2d1 --- /dev/null +++ b/devito/petsc/iet/__init__.py @@ -0,0 +1 @@ +from devito.petsc.iet import * # noqa diff --git a/devito/petsc/iet/nodes.py b/devito/petsc/iet/nodes.py new file mode 100644 index 0000000000..2137495487 --- /dev/null +++ b/devito/petsc/iet/nodes.py @@ -0,0 +1,39 @@ +from devito.ir.iet import Expression, Callback, FixedArgsCallable, Call +from devito.ir.equations import OpInjectSolve + + +class LinearSolverExpression(Expression): + """ + Base class for general expressions required by a + matrix-free linear solve of the form Ax=b. + """ + pass + + +class InjectSolveDummy(LinearSolverExpression): + """ + Placeholder expression to run the iterative solver. + """ + def __init__(self, expr, pragmas=None, operation=OpInjectSolve): + super().__init__(expr, pragmas=pragmas, operation=operation) + + +class PETScCallable(FixedArgsCallable): + pass + + +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 + + +class PETScCall(Call): + pass diff --git a/devito/petsc/iet/passes.py b/devito/petsc/iet/passes.py new file mode 100644 index 0000000000..561f9fd0f4 --- /dev/null +++ b/devito/petsc/iet/passes.py @@ -0,0 +1,136 @@ +import cgen as c + +from devito.passes.iet.engine import iet_pass +from devito.ir.iet import Transformer, MapNodes, Iteration, BlankLine +from devito.symbolics import Byref, Macro +from devito.petsc.types import (PetscMPIInt, PetscErrorCode) +from devito.petsc.iet.nodes import InjectSolveDummy +from devito.petsc.utils import core_metadata +from devito.petsc.iet.routines import (CallbackBuilder, BaseObjectBuilder, BaseSetup, + Solver, TimeDependent, NonTimeDependent) +from devito.petsc.iet.utils import petsc_call, petsc_call_mpi + + +@iet_pass +def lower_petsc(iet, **kwargs): + # Check if PETScSolve was used + injectsolve_mapper = MapNodes(Iteration, InjectSolveDummy, + 'groupby').visit(iet) + + if not injectsolve_mapper: + return iet, {} + + targets = [i.expr.rhs.target for (i,) in injectsolve_mapper.values()] + init = init_petsc(**kwargs) + + # Assumption is that all targets have the same grid so can use any target here + objs = build_core_objects(targets[-1], **kwargs) + + # Create core PETSc calls (not specific to each PETScSolve) + core = make_core_petsc_calls(objs, **kwargs) + + setup = [] + subs = {} + efuncs = {} + + for iters, (injectsolve,) in injectsolve_mapper.items(): + + builder = Builder(injectsolve, objs, iters, **kwargs) + + setup.extend(builder.solversetup.calls) + + # Transform the spatial iteration loop with the calls to execute the solver + subs.update(builder.solve.mapper) + + efuncs.update(builder.cbbuilder.efuncs) + + iet = Transformer(subs).visit(iet) + + body = core + tuple(setup) + (BlankLine,) + iet.body.body + body = iet.body._rebuild( + init=init, body=body, + frees=(petsc_call('PetscFinalize', []),) + ) + iet = iet._rebuild(body=body) + metadata = core_metadata() + metadata.update({'efuncs': tuple(efuncs.values())}) + + return iet, metadata + + +def init_petsc(**kwargs): + # Initialize PETSc -> for now, assuming all solver options have to be + # specified via the parameters dict in PETScSolve + # TODO: Are users going to be able to use PETSc command line arguments? + # In firedrake, they have an options_prefix for each solver, enabling the use + # of command line options + initialize = petsc_call('PetscInitialize', [Null, Null, Null, Null]) + + return petsc_func_begin_user, initialize + + +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 + } + + +class Builder: + """ + This class is designed to support future extensions, enabling + different combinations of solver types, preconditioning methods, + and other functionalities as needed. + + The class will be extended to accommodate different solver types by + returning subclasses of the objects initialised in __init__, + depending on the properties of `injectsolve`. + """ + def __init__(self, injectsolve, objs, iters, **kwargs): + + # Determine the time dependency class + time_mapper = injectsolve.expr.rhs.time_mapper + timedep = TimeDependent if time_mapper else NonTimeDependent + self.timedep = timedep(injectsolve, iters, **kwargs) + + # Objects + self.objbuilder = BaseObjectBuilder(injectsolve, **kwargs) + self.solver_objs = self.objbuilder.solver_objs + + # Callbacks + self.cbbuilder = CallbackBuilder( + injectsolve, objs, self.solver_objs, timedep=self.timedep, + **kwargs + ) + + # Solver setup + self.solversetup = BaseSetup( + self.solver_objs, objs, injectsolve, self.cbbuilder + ) + + # Execute the solver + self.solve = Solver( + self.solver_objs, objs, injectsolve, iters, + self.cbbuilder, timedep=self.timedep + ) + + +Null = Macro('NULL') +void = 'void' + + +# TODO: Don't use c.Line here? +petsc_func_begin_user = c.Line('PetscFunctionBeginUser;') diff --git a/devito/petsc/iet/routines.py b/devito/petsc/iet/routines.py new file mode 100644 index 0000000000..81547ac64b --- /dev/null +++ b/devito/petsc/iet/routines.py @@ -0,0 +1,1003 @@ +from collections import OrderedDict + +import cgen as c + +from devito.ir.iet import (Call, FindSymbols, List, Uxreplace, CallableBody, + Dereference, DummyExpr, BlankLine, Callable, FindNodes, + retrieve_iteration_tree, filter_iterations) +from devito.symbolics import (Byref, FieldFromPointer, Macro, cast_mapper, + FieldFromComposite) +from devito.symbolics.unevaluation import Mul +from devito.types.basic import AbstractFunction +from devito.types import Temp, Symbol +from devito.tools import filter_ordered + +from devito.petsc.types import PETScArray +from devito.petsc.iet.nodes import (PETScCallable, FormFunctionCallback, + MatVecCallback, InjectSolveDummy) +from devito.petsc.iet.utils import petsc_call, petsc_struct +from devito.petsc.utils import solver_mapper +from devito.petsc.types import (DM, CallbackDM, Mat, LocalVec, GlobalVec, KSP, PC, + SNES, DummyArg, PetscInt, StartPtr) + + +class CallbackBuilder: + """ + Build IET routines to generate PETSc callback functions. + """ + def __init__(self, injectsolve, objs, solver_objs, + rcompile=None, sregistry=None, timedep=None, **kwargs): + + self.rcompile = rcompile + self.sregistry = sregistry + self.timedep = timedep + self.solver_objs = solver_objs + + self._efuncs = OrderedDict() + self._struct_params = [] + + self._matvec_callback = None + self._formfunc_callback = None + self._formrhs_callback = None + self._struct_callback = None + + self._make_core(injectsolve, objs, solver_objs) + self._main_struct(solver_objs) + self._make_struct_callback(solver_objs, objs) + self._local_struct(solver_objs) + self._efuncs = self._uxreplace_efuncs() + + @property + def efuncs(self): + return self._efuncs + + @property + def struct_params(self): + return self._struct_params + + @property + def filtered_struct_params(self): + return filter_ordered(self.struct_params) + + @property + def matvec_callback(self): + return self._matvec_callback + + @property + def formfunc_callback(self): + return self._formfunc_callback + + @property + def formrhs_callback(self): + return self._formrhs_callback + + @property + def struct_callback(self): + return self._struct_callback + + def _make_core(self, injectsolve, objs, solver_objs): + self._make_matvec(injectsolve, objs, solver_objs) + self._make_formfunc(injectsolve, objs, solver_objs) + self._make_formrhs(injectsolve, objs, solver_objs) + + def _make_matvec(self, injectsolve, objs, solver_objs): + # Compile matvec `eqns` into an IET via recursive compilation + irs_matvec, _ = self.rcompile(injectsolve.expr.rhs.matvecs, + options={'mpi': False}, sregistry=self.sregistry) + body_matvec = self._create_matvec_body(injectsolve, + List(body=irs_matvec.uiet.body), + solver_objs, objs) + + matvec_callback = PETScCallable( + self.sregistry.make_name(prefix='MyMatShellMult_'), body_matvec, + retval=objs['err'], + parameters=( + solver_objs['Jac'], solver_objs['X_global'], solver_objs['Y_global'] + ) + ) + self._matvec_callback = matvec_callback + self._efuncs[matvec_callback.name] = matvec_callback + + def _create_matvec_body(self, injectsolve, body, solver_objs, objs): + linsolve_expr = injectsolve.expr.rhs + + dmda = solver_objs['callbackdm'] + + body = self.timedep.uxreplace_time(body) + + fields = self._dummy_fields(body, solver_objs) + + y_matvec = linsolve_expr.arrays['y_matvec'] + x_matvec = linsolve_expr.arrays['x_matvec'] + + mat_get_dm = petsc_call('MatGetDM', [solver_objs['Jac'], Byref(dmda)]) + + dm_get_app_context = petsc_call( + 'DMGetApplicationContext', [dmda, Byref(dummyctx._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(linsolve_expr.localinfo)] + ) + + 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'] + ]) + + dm_restore_local_xvec = petsc_call( + 'DMRestoreLocalVector', [dmda, Byref(solver_objs['X_local'])] + ) + + dm_restore_local_yvec = petsc_call( + 'DMRestoreLocalVector', [dmda, Byref(solver_objs['Y_local'])] + ) + + # TODO: Some of the calls are placed in the `stacks` argument of the + # `CallableBody` to ensure that they precede the `cast` statements. The + # 'casts' depend on the calls, so this order is necessary. By doing this, + # 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, + dm_restore_local_xvec, + dm_restore_local_yvec) + ) + + 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 + ) + + # Dereference function data in struct + dereference_funcs = [Dereference(i, dummyctx) for i in + fields if isinstance(i.function, AbstractFunction)] + + matvec_body = CallableBody( + List(body=body), + init=(petsc_func_begin_user,), + stacks=stacks+tuple(dereference_funcs), + retstmt=(Call('PetscFunctionReturn', arguments=[0]),) + ) + + # Replace non-function data with pointer to data in struct + subs = {i._C_symbol: FieldFromPointer(i._C_symbol, dummyctx) for i in fields} + matvec_body = Uxreplace(subs).visit(matvec_body) + + self._struct_params.extend(fields) + + return matvec_body + + def _make_formfunc(self, injectsolve, objs, solver_objs): + # Compile formfunc `eqns` into an IET via recursive compilation + irs_formfunc, _ = self.rcompile( + injectsolve.expr.rhs.formfuncs, + options={'mpi': False}, sregistry=self.sregistry + ) + body_formfunc = self._create_formfunc_body(injectsolve, + List(body=irs_formfunc.uiet.body), + solver_objs, objs) + + formfunc_callback = PETScCallable( + self.sregistry.make_name(prefix='FormFunction_'), body_formfunc, + retval=objs['err'], + parameters=(solver_objs['snes'], solver_objs['X_global'], + solver_objs['F_global'], dummyptr) + ) + self._formfunc_callback = formfunc_callback + self._efuncs[formfunc_callback.name] = formfunc_callback + + def _create_formfunc_body(self, injectsolve, body, solver_objs, objs): + linsolve_expr = injectsolve.expr.rhs + + dmda = solver_objs['callbackdm'] + + body = self.timedep.uxreplace_time(body) + + fields = self._dummy_fields(body, solver_objs) + + f_formfunc = linsolve_expr.arrays['f_formfunc'] + x_formfunc = linsolve_expr.arrays['x_formfunc'] + + snes_get_dm = petsc_call('SNESGetDM', [solver_objs['snes'], Byref(dmda)]) + + dm_get_app_context = petsc_call( + 'DMGetApplicationContext', [dmda, Byref(dummyctx._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['F_local'])] + ) + + vec_get_array_y = petsc_call( + 'VecGetArray', [solver_objs['F_local'], Byref(f_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(linsolve_expr.localinfo)] + ) + + vec_restore_array_y = petsc_call( + 'VecRestoreArray', [solver_objs['F_local'], Byref(f_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['F_local'], 'INSERT_VALUES', solver_objs['F_global'] + ]) + + dm_local_to_global_end = petsc_call('DMLocalToGlobalEnd', [ + dmda, solver_objs['F_local'], 'INSERT_VALUES', solver_objs['F_global'] + ]) + + dm_restore_local_xvec = petsc_call( + 'DMRestoreLocalVector', [dmda, Byref(solver_objs['X_local'])] + ) + + dm_restore_local_yvec = petsc_call( + 'DMRestoreLocalVector', [dmda, Byref(solver_objs['F_local'])] + ) + + body = body._rebuild( + body=body.body + + (vec_restore_array_y, + vec_restore_array_x, + dm_local_to_global_begin, + dm_local_to_global_end, + dm_restore_local_xvec, + dm_restore_local_yvec) + ) + + 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 + ) + + # Dereference function data in struct + dereference_funcs = [Dereference(i, dummyctx) for i in + fields if isinstance(i.function, AbstractFunction)] + + formfunc_body = CallableBody( + List(body=body), + init=(petsc_func_begin_user,), + stacks=stacks+tuple(dereference_funcs), + retstmt=(Call('PetscFunctionReturn', arguments=[0]),)) + + # Replace non-function data with pointer to data in struct + subs = {i._C_symbol: FieldFromPointer(i._C_symbol, dummyctx) for i in fields} + formfunc_body = Uxreplace(subs).visit(formfunc_body) + + self._struct_params.extend(fields) + + return formfunc_body + + def _make_formrhs(self, injectsolve, objs, solver_objs): + # Compile formrhs `eqns` into an IET via recursive compilation + irs_formrhs, _ = self.rcompile(injectsolve.expr.rhs.formrhs, + options={'mpi': False}, sregistry=self.sregistry) + body_formrhs = self._create_form_rhs_body(injectsolve, + List(body=irs_formrhs.uiet.body), + solver_objs, objs) + + formrhs_callback = PETScCallable( + self.sregistry.make_name(prefix='FormRHS_'), body_formrhs, retval=objs['err'], + parameters=( + solver_objs['snes'], solver_objs['b_local'] + ) + ) + self._formrhs_callback = formrhs_callback + self._efuncs[formrhs_callback.name] = formrhs_callback + + def _create_form_rhs_body(self, injectsolve, body, solver_objs, objs): + linsolve_expr = injectsolve.expr.rhs + + dmda = solver_objs['callbackdm'] + + snes_get_dm = petsc_call('SNESGetDM', [solver_objs['snes'], Byref(dmda)]) + + b_arr = linsolve_expr.arrays['b_tmp'] + + vec_get_array = petsc_call( + 'VecGetArray', [solver_objs['b_local'], Byref(b_arr._C_symbol)] + ) + + dm_get_local_info = petsc_call( + 'DMDAGetLocalInfo', [dmda, Byref(linsolve_expr.localinfo)] + ) + + body = self.timedep.uxreplace_time(body) + + fields = self._dummy_fields(body, solver_objs) + + dm_get_app_context = petsc_call( + 'DMGetApplicationContext', [dmda, Byref(dummyctx._C_symbol)] + ) + + vec_restore_array = petsc_call( + 'VecRestoreArray', [solver_objs['b_local'], Byref(b_arr._C_symbol)] + ) + + body = body._rebuild(body=body.body + (vec_restore_array,)) + + stacks = ( + snes_get_dm, + dm_get_app_context, + vec_get_array, + dm_get_local_info + ) + + # Dereference function data in struct + dereference_funcs = [Dereference(i, dummyctx) for i in + fields if isinstance(i.function, AbstractFunction)] + + formrhs_body = CallableBody( + List(body=[body]), + init=(petsc_func_begin_user,), + stacks=stacks+tuple(dereference_funcs), + retstmt=(Call('PetscFunctionReturn', arguments=[0]),) + ) + + # Replace non-function data with pointer to data in struct + subs = {i._C_symbol: FieldFromPointer(i._C_symbol, dummyctx) for + i in fields if not isinstance(i.function, AbstractFunction)} + formrhs_body = Uxreplace(subs).visit(formrhs_body) + + self._struct_params.extend(fields) + + return formrhs_body + + def _local_struct(self, solver_objs): + """ + This is the struct used within callback functions, + usually accessed via DMGetApplicationContext. + """ + solver_objs['localctx'] = petsc_struct( + dummyctx.name, + self.filtered_struct_params, + solver_objs['Jac'].name+'_ctx', + liveness='eager' + ) + + def _main_struct(self, solver_objs): + """ + This is the struct initialised inside the main kernel and + attached to the DM via DMSetApplicationContext. + """ + solver_objs['mainctx'] = petsc_struct( + self.sregistry.make_name(prefix='ctx'), + self.filtered_struct_params, + solver_objs['Jac'].name+'_ctx' + ) + + def _make_struct_callback(self, solver_objs, objs): + mainctx = solver_objs['mainctx'] + body = [ + DummyExpr(FieldFromPointer(i._C_symbol, mainctx), i._C_symbol) + for i in mainctx.callback_fields + ] + struct_callback_body = CallableBody( + List(body=body), init=(petsc_func_begin_user,), + retstmt=tuple([Call('PetscFunctionReturn', arguments=[0])]) + ) + struct_callback = Callable( + self.sregistry.make_name(prefix='PopulateMatContext_'), + struct_callback_body, objs['err'], + parameters=[mainctx] + ) + self._efuncs[struct_callback.name] = struct_callback + self._struct_callback = struct_callback + + def _dummy_fields(self, iet, solver_objs): + # Place all context data required by the shell routines into a struct + fields = [f.function for f in FindSymbols('basics').visit(iet)] + fields = [f for f in fields if not isinstance(f.function, (PETScArray, Temp))] + fields = [ + f for f in fields if not (f.is_Dimension and not (f.is_Time or f.is_Modulo)) + ] + return fields + + def _uxreplace_efuncs(self): + mapper = {} + visitor = Uxreplace({dummyctx: self.solver_objs['localctx']}) + for k, v in self._efuncs.items(): + mapper.update({k: visitor.visit(v)}) + return mapper + + +class BaseObjectBuilder: + """ + A base class for constructing objects needed for a PETSc solver. + Designed to be extended by subclasses, which can override the `_extend_build` + method to support specific use cases. + """ + + def __init__(self, injectsolve, sregistry=None, **kwargs): + self.sregistry = sregistry + self.solver_objs = self._build(injectsolve) + + def _build(self, injectsolve): + """ + Constructs the core dictionary of solver objects and allows + subclasses to extend or modify it via `_extend_build`. + + Returns: + dict: A dictionary containing the following objects: + - 'Jac' (Mat): A matrix representing the jacobian. + - 'x_global' (GlobalVec): The global solution vector. + - 'x_local' (LocalVec): The local solution vector. + - 'b_global': (GlobalVec) Global RHS vector `b`, where `F(x) = b`. + - 'b_local': (LocalVec) Local RHS vector `b`, where `F(x) = b`. + - 'ksp': (KSP) Krylov solver object that manages the linear solver. + - 'pc': (PC) Preconditioner object. + - 'snes': (SNES) Nonlinear solver object. + - 'F_global': (GlobalVec) Global residual vector `F`, where `F(x) = b`. + - 'F_local': (LocalVec) Local residual vector `F`, where `F(x) = b`. + - 'Y_global': (GlobalVector) The output vector populated by the + matrix-free `MyMatShellMult` callback function. + - 'Y_local': (LocalVector) The output vector populated by the matrix-free + `MyMatShellMult` callback function. + - 'X_global': (GlobalVec) Current guess for the solution, + required by the FormFunction callback. + - 'X_local': (LocalVec) Current guess for the solution, + required by the FormFunction callback. + - 'localsize' (PetscInt): The local length of the solution vector. + - 'start_ptr' (StartPtr): A pointer to the beginning of the solution array + that will be updated at each time step. + - 'dmda' (DM): The DMDA object associated with this solve, linked to + the SNES object via `SNESSetDM`. + - 'callbackdm' (CallbackDM): The DM object accessed within callback + functions via `SNESGetDM`. + """ + target = injectsolve.expr.rhs.target + sreg = self.sregistry + base_dict = { + 'Jac': Mat(sreg.make_name(prefix='J_')), + 'x_global': GlobalVec(sreg.make_name(prefix='x_global_')), + 'x_local': LocalVec(sreg.make_name(prefix='x_local_'), liveness='eager'), + 'b_global': GlobalVec(sreg.make_name(prefix='b_global_')), + 'b_local': LocalVec(sreg.make_name(prefix='b_local_')), + 'ksp': KSP(sreg.make_name(prefix='ksp_')), + 'pc': PC(sreg.make_name(prefix='pc_')), + 'snes': SNES(sreg.make_name(prefix='snes_')), + 'F_global': GlobalVec(sreg.make_name(prefix='F_global_')), + 'F_local': LocalVec(sreg.make_name(prefix='F_local_'), liveness='eager'), + 'Y_global': GlobalVec(sreg.make_name(prefix='Y_global_')), + 'Y_local': LocalVec(sreg.make_name(prefix='Y_local_'), liveness='eager'), + 'X_global': GlobalVec(sreg.make_name(prefix='X_global_')), + 'X_local': LocalVec(sreg.make_name(prefix='X_local_'), liveness='eager'), + 'localsize': PetscInt(sreg.make_name(prefix='localsize_')), + 'start_ptr': StartPtr(sreg.make_name(prefix='start_ptr_'), target.dtype), + 'dmda': DM(sreg.make_name(prefix='da_'), liveness='eager', + stencil_width=target.space_order), + 'callbackdm': CallbackDM(sreg.make_name(prefix='dm_'), + liveness='eager', stencil_width=target.space_order), + } + return self._extend_build(base_dict, injectsolve) + + def _extend_build(self, base_dict, injectsolve): + """ + Subclasses can override this method to extend or modify the + base dictionary of solver objects. + """ + return base_dict + + +class BaseSetup: + def __init__(self, solver_objs, objs, injectsolve, cbbuilder): + self.calls = self._setup(solver_objs, objs, injectsolve, cbbuilder) + + def _setup(self, solver_objs, objs, injectsolve, cbbuilder): + dmda = solver_objs['dmda'] + + solver_params = injectsolve.expr.rhs.solver_parameters + + snes_create = petsc_call('SNESCreate', [objs['comm'], Byref(solver_objs['snes'])]) + + snes_set_dm = petsc_call('SNESSetDM', [solver_objs['snes'], dmda]) + + create_matrix = petsc_call('DMCreateMatrix', [dmda, Byref(solver_objs['Jac'])]) + + # NOTE: Assuming all solves are linear for now. + snes_set_type = petsc_call('SNESSetType', [solver_objs['snes'], 'SNESKSPONLY']) + + snes_set_jac = petsc_call( + 'SNESSetJacobian', [solver_objs['snes'], solver_objs['Jac'], + solver_objs['Jac'], 'MatMFFDComputeJacobian', Null] + ) + + global_x = petsc_call('DMCreateGlobalVector', + [dmda, Byref(solver_objs['x_global'])]) + + global_b = petsc_call('DMCreateGlobalVector', + [dmda, Byref(solver_objs['b_global'])]) + + local_b = petsc_call('DMCreateLocalVector', + [dmda, Byref(solver_objs['b_local'])]) + + snes_get_ksp = petsc_call('SNESGetKSP', + [solver_objs['snes'], Byref(solver_objs['ksp'])]) + + ksp_set_tols = petsc_call( + 'KSPSetTolerances', [solver_objs['ksp'], solver_params['ksp_rtol'], + solver_params['ksp_atol'], solver_params['ksp_divtol'], + solver_params['ksp_max_it']] + ) + + ksp_set_type = petsc_call( + 'KSPSetType', [solver_objs['ksp'], solver_mapper[solver_params['ksp_type']]] + ) + + ksp_get_pc = petsc_call( + 'KSPGetPC', [solver_objs['ksp'], Byref(solver_objs['pc'])] + ) + + # Even though the default will be jacobi, set to PCNONE for now + pc_set_type = petsc_call('PCSetType', [solver_objs['pc'], 'PCNONE']) + + ksp_set_from_ops = petsc_call('KSPSetFromOptions', [solver_objs['ksp']]) + + matvec_operation = petsc_call( + 'MatShellSetOperation', + [solver_objs['Jac'], 'MATOP_MULT', + MatVecCallback(cbbuilder.matvec_callback.name, void, void)] + ) + + formfunc_operation = petsc_call( + 'SNESSetFunction', + [solver_objs['snes'], Null, + FormFunctionCallback(cbbuilder.formfunc_callback.name, void, void), Null] + ) + + dmda_calls = self._create_dmda_calls(dmda, objs) + + mainctx = solver_objs['mainctx'] + + call_struct_callback = petsc_call( + cbbuilder.struct_callback.name, [Byref(mainctx)] + ) + calls_set_app_ctx = [ + petsc_call('DMSetApplicationContext', [dmda, Byref(mainctx)]) + ] + calls = [call_struct_callback] + calls_set_app_ctx + [BlankLine] + + base_setup = dmda_calls + ( + snes_create, + snes_set_dm, + create_matrix, + snes_set_jac, + snes_set_type, + global_x, + global_b, + local_b, + snes_get_ksp, + ksp_set_tols, + ksp_set_type, + ksp_get_pc, + pc_set_type, + ksp_set_from_ops, + matvec_operation, + formfunc_operation, + ) + tuple(calls) + + extended_setup = self._extend_setup(solver_objs, objs, injectsolve, cbbuilder) + return base_setup + tuple(extended_setup) + + def _extend_setup(self, solver_objs, objs, injectsolve, cbbuilder): + """ + Hook for subclasses to add additional setup calls. + """ + return [] + + def _create_dmda_calls(self, dmda, objs): + dmda_create = self._create_dmda(dmda, objs) + dm_setup = petsc_call('DMSetUp', [dmda]) + dm_mat_type = petsc_call('DMSetMatType', [dmda, 'MATSHELL']) + return dmda_create, dm_setup, dm_mat_type + + def _create_dmda(self, dmda, objs): + grid = objs['grid'] + + nspace_dims = len(grid.dimensions) + + # MPI communicator + args = [objs['comm']] + + # Type of ghost nodes + args.extend(['DM_BOUNDARY_GHOSTED' for _ in range(nspace_dims)]) + + # Stencil type + if nspace_dims > 1: + args.append('DMDA_STENCIL_BOX') + + # Global dimensions + args.extend(list(grid.shape)[::-1]) + # No.of processors in each dimension + if nspace_dims > 1: + args.extend(list(grid.distributor.topology)[::-1]) + + # Number of degrees of freedom per node + args.append(1) + # "Stencil width" -> size of overlap + args.append(dmda.stencil_width) + args.extend([Null]*nspace_dims) + + # The distributed array object + args.append(Byref(dmda)) + + # The PETSc call used to create the DMDA + dmda = petsc_call('DMDACreate%sd' % nspace_dims, args) + + return dmda + + +class Solver: + def __init__(self, solver_objs, objs, injectsolve, iters, cbbuilder, + timedep=None, **kwargs): + self.timedep = timedep + self.calls = self._execute_solve(solver_objs, objs, injectsolve, iters, cbbuilder) + self.spatial_body = self._spatial_loop_nest(iters, injectsolve) + + space_iter, = self.spatial_body + self.mapper = {space_iter: self.calls} + + def _execute_solve(self, solver_objs, objs, injectsolve, iters, cbbuilder): + """ + Assigns the required time iterators to the struct and executes + the necessary calls to execute the SNES solver. + """ + struct_assignment = self.timedep.assign_time_iters(solver_objs['mainctx']) + + rhs_callback = cbbuilder.formrhs_callback + + dmda = solver_objs['dmda'] + + rhs_call = petsc_call(rhs_callback.name, list(rhs_callback.parameters)) + + local_x = petsc_call('DMCreateLocalVector', + [dmda, Byref(solver_objs['x_local'])]) + + vec_replace_array = self.timedep.replace_array(solver_objs) + + 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']] + ) + + run_solver_calls = (struct_assignment,) + ( + 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, + ) + return List(body=run_solver_calls) + + def _spatial_loop_nest(self, iters, injectsolve): + spatial_body = [] + for tree in retrieve_iteration_tree(iters[0]): + root = filter_iterations(tree, key=lambda i: i.dim.is_Space)[0] + if injectsolve in FindNodes(InjectSolveDummy).visit(root): + spatial_body.append(root) + return spatial_body + + +class NonTimeDependent: + def __init__(self, injectsolve, iters, **kwargs): + self.injectsolve = injectsolve + self.iters = iters + self.kwargs = kwargs + self.origin_to_moddim = self._origin_to_moddim_mapper(iters) + self.time_idx_to_symb = injectsolve.expr.rhs.time_mapper + + @property + def is_target_time(self): + return False + + @property + def target(self): + return self.injectsolve.expr.rhs.target + + def _origin_to_moddim_mapper(self, iters): + return {} + + def uxreplace_time(self, body): + return body + + def replace_array(self, solver_objs): + """ + VecReplaceArray() is a PETSc function that allows replacing the array + of a `Vec` with a user provided array. + https://petsc.org/release/manualpages/Vec/VecReplaceArray/ + + This function is used to replace the array of the PETSc solution `Vec` + with the array from the `Function` object representing the target. + + Examples + -------- + >>> self.target + f1(x, y) + >>> call = replace_array(solver_objs) + >>> print(call) + PetscCall(VecReplaceArray(x_local_0,f1_vec->data)); + """ + field_from_ptr = FieldFromPointer( + self.target.function._C_field_data, self.target.function._C_symbol + ) + vec_replace_array = (petsc_call( + 'VecReplaceArray', [solver_objs['x_local'], field_from_ptr] + ),) + return vec_replace_array + + def assign_time_iters(self, struct): + return [] + + +class TimeDependent(NonTimeDependent): + """ + A class for managing time-dependent solvers. + + This includes scenarios where the target is not directly a `TimeFunction`, + but depends on other functions that are. + + Outline of time loop abstraction with PETSc: + + - At PETScSolve, time indices are replaced with temporary `Symbol` objects + via a mapper (e.g., {t: tau0, t + dt: tau1}) to prevent the time loop + from being generated in the callback functions. These callbacks, needed + for each `SNESSolve` at every time step, don't require the time loop, but + may still need access to data from other time steps. + - All `Function` objects are passed through the initial lowering via the + `LinearSolveExpr` object, ensuring the correct time loop is generated + in the main kernel. + - Another mapper is created based on the modulo dimensions + generated by the `LinearSolveExpr` object in the main kernel + (e.g., {time: time, t: t0, t + 1: t1}). + - These two mappers are used to generate a final mapper `symb_to_moddim` + (e.g. {tau0: t0, tau1: t1}) which is used at the IET level to + replace the temporary `Symbol` objects in the callback functions with + the correct modulo dimensions. + - Modulo dimensions are updated in the matrix context struct at each time + step and can be accessed in the callback functions where needed. + """ + @property + def is_target_time(self): + return any(i.is_Time for i in self.target.dimensions) + + @property + def time_spacing(self): + return self.target.grid.stepping_dim.spacing + + @property + def target_time(self): + target_time = [ + i for i, d in zip(self.target.indices, self.target.dimensions) + if d.is_Time + ] + assert len(target_time) == 1 + target_time = target_time.pop() + return target_time + + @property + def symb_to_moddim(self): + """ + Maps temporary `Symbol` objects created during `PETScSolve` to their + corresponding modulo dimensions (e.g. creates {tau0: t0, tau1: t1}). + """ + mapper = { + v: k.xreplace({self.time_spacing: 1, -self.time_spacing: -1}) + for k, v in self.time_idx_to_symb.items() + } + return {symb: self.origin_to_moddim[mapper[symb]] for symb in mapper} + + def uxreplace_time(self, body): + return Uxreplace(self.symb_to_moddim).visit(body) + + def _origin_to_moddim_mapper(self, iters): + """ + Creates a mapper of the origin of the time dimensions to their corresponding + modulo dimensions from a list of `Iteration` objects. + + Examples + -------- + >>> iters + (, + ) + >>> _origin_to_moddim_mapper(iters) + {time: time, t: t0, t + 1: t1} + """ + time_iter = [i for i in iters if any(d.is_Time for d in i.dimensions)] + mapper = {} + + if not time_iter: + return mapper + + for i in time_iter: + for d in i.dimensions: + if d.is_Modulo: + mapper[d.origin] = d + elif d.is_Time: + mapper[d] = d + return mapper + + def replace_array(self, solver_objs): + """ + In the case that the actual target is time-dependent e.g a `TimeFunction`, + a pointer to the first element in the array that will be updated during + the time step is passed to VecReplaceArray(). + + Examples + -------- + >>> self.target + f1(time + dt, x, y) + >>> calls = replace_array(solver_objs) + >>> print(List(body=calls)) + PetscCall(VecGetSize(x_local_0,&(localsize_0))); + float * start_ptr_0 = (time + 1)*localsize_0 + (float*)(f1_vec->data); + PetscCall(VecReplaceArray(x_local_0,start_ptr_0)); + + >>> self.target + f1(t + dt, x, y) + >>> calls = replace_array(solver_objs) + >>> print(List(body=calls)) + PetscCall(VecGetSize(x_local_0,&(localsize_0))); + float * start_ptr_0 = t1*localsize_0 + (float*)(f1_vec->data); + """ + if self.is_target_time: + mapper = {self.time_spacing: 1, -self.time_spacing: -1} + target_time = self.target_time.xreplace(mapper) + + try: + target_time = self.origin_to_moddim[target_time] + except KeyError: + pass + + start_ptr = solver_objs['start_ptr'] + + vec_get_size = petsc_call( + 'VecGetSize', [solver_objs['x_local'], Byref(solver_objs['localsize'])] + ) + + field_from_ptr = FieldFromPointer( + self.target.function._C_field_data, self.target.function._C_symbol + ) + + expr = DummyExpr( + start_ptr, cast_mapper[(self.target.dtype, '*')](field_from_ptr) + + Mul(target_time, solver_objs['localsize']), init=True + ) + + vec_replace_array = petsc_call( + 'VecReplaceArray', [solver_objs['x_local'], start_ptr] + ) + return (vec_get_size, expr, vec_replace_array) + else: + return super().replace_array(solver_objs) + + def assign_time_iters(self, struct): + """ + Assign required time iterators to the struct. + These iterators are updated at each timestep in the main kernel + for use in callback functions. + + Examples + -------- + >>> struct + ctx + >>> struct.fields + [h_x, x_M, x_m, f1(t, x), t0, t1] + >>> assigned = assign_time_iters(struct) + >>> print(assigned[0]) + ctx.t0 = t0; + >>> print(assigned[1]) + ctx.t1 = t1; + """ + to_assign = [ + f for f in struct.fields if (f.is_Dimension and (f.is_Time or f.is_Modulo)) + ] + time_iter_assignments = [ + DummyExpr(FieldFromComposite(field, struct), field) + for field in to_assign + ] + return time_iter_assignments + + +Null = Macro('NULL') +void = 'void' +dummyctx = Symbol('lctx') +dummyptr = DummyArg('dummy') + + +# TODO: Don't use c.Line here? +petsc_func_begin_user = c.Line('PetscFunctionBeginUser;') diff --git a/devito/petsc/iet/utils.py b/devito/petsc/iet/utils.py new file mode 100644 index 0000000000..adcf709eab --- /dev/null +++ b/devito/petsc/iet/utils.py @@ -0,0 +1,22 @@ +from devito.petsc.iet.nodes import InjectSolveDummy, PETScCall +from devito.ir.equations import OpInjectSolve + + +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, pname, liveness='lazy'): + # TODO: Fix this circular import + from devito.petsc.types.object import PETScStruct + return PETScStruct(name=name, pname=pname, + fields=fields, liveness=liveness) + + +# 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} diff --git a/devito/petsc/solve.py b/devito/petsc/solve.py new file mode 100644 index 0000000000..4f16ded1f3 --- /dev/null +++ b/devito/petsc/solve.py @@ -0,0 +1,241 @@ +from functools import singledispatch + +import sympy + +from devito.finite_differences.differentiable import Mul +from devito.finite_differences.derivative import Derivative +from devito.types import Eq, Symbol, SteppingDimension, TimeFunction +from devito.types.equation import InjectSolveEq +from devito.operations.solve import eval_time_derivatives +from devito.symbolics import retrieve_functions +from devito.tools import as_tuple +from devito.petsc.types import LinearSolveExpr, PETScArray, DMDALocalInfo + + +__all__ = ['PETScSolve'] + + +def PETScSolve(eqns, target, solver_parameters=None, **kwargs): + prefixes = ['y_matvec', 'x_matvec', 'f_formfunc', 'x_formfunc', 'b_tmp'] + + localinfo = DMDALocalInfo(name='info', liveness='eager') + + arrays = { + p: PETScArray(name='%s_%s' % (p, target.name), + target=target, + liveness='eager', + localinfo=localinfo) + for p in prefixes + } + + matvecs = [] + formfuncs = [] + formrhs = [] + + eqns = as_tuple(eqns) + + for eq in eqns: + b, F_target, targets = separate_eqn(eq, target) + + # TODO: Current assumption is that problem is linear and user has not provided + # a jacobian. Hence, we can use F_target to form the jac-vec product + matvecs.append(Eq( + arrays['y_matvec'], + F_target.subs(targets_to_arrays(arrays['x_matvec'], targets)), + subdomain=eq.subdomain + )) + + formfuncs.append(Eq( + arrays['f_formfunc'], + F_target.subs(targets_to_arrays(arrays['x_formfunc'], targets)), + subdomain=eq.subdomain + )) + + formrhs.append(Eq( + arrays['b_tmp'], + b, + subdomain=eq.subdomain + )) + + funcs = retrieve_functions(eqns) + time_mapper = generate_time_mapper(funcs) + + matvecs, formfuncs, formrhs = ( + [eq.xreplace(time_mapper) for eq in lst] for lst in (matvecs, formfuncs, formrhs) + ) + # Placeholder equation for inserting calls to the solver and generating + # correct time loop etc + inject_solve = InjectSolveEq(target, LinearSolveExpr( + expr=tuple(funcs), + target=target, + solver_parameters=solver_parameters, + matvecs=matvecs, + formfuncs=formfuncs, + formrhs=formrhs, + arrays=arrays, + time_mapper=time_mapper, + localinfo=localinfo + )) + + return [inject_solve] + + +def separate_eqn(eqn, target): + """ + Separate the equation into two separate expressions, + where F(target) = b. + """ + zeroed_eqn = Eq(eqn.lhs - eqn.rhs, 0) + zeroed_eqn = eval_time_derivatives(zeroed_eqn.lhs) + target_funcs = set(generate_targets(zeroed_eqn, target)) + b, F_target = remove_targets(zeroed_eqn, target_funcs) + return -b, F_target, target_funcs + + +def generate_targets(eq, target): + """ + Extract all the functions that share the same time index as the target + but may have different spatial indices. + """ + funcs = retrieve_functions(eq) + if isinstance(target, TimeFunction): + time_idx = target.indices[target.time_dim] + targets = [ + f for f in funcs if f.function is target.function and time_idx + in f.indices + ] + else: + targets = [f for f in funcs if f.function is target.function] + return targets + + +def targets_to_arrays(array, targets): + """ + Map each target in `targets` to a corresponding array generated from `array`, + matching the spatial indices of the target. + + Example: + -------- + >>> array + vec_u(x, y) + + >>> targets + {u(t + dt, x + h_x, y), u(t + dt, x - h_x, y), u(t + dt, x, y)} + + >>> targets_to_arrays(array, targets) + {u(t + dt, x - h_x, y): vec_u(x - h_x, y), + u(t + dt, x + h_x, y): vec_u(x + h_x, y), + u(t + dt, x, y): vec_u(x, y)} + """ + space_indices = [ + tuple(f.indices[d] for d in f.space_dimensions) for f in targets + ] + array_targets = [ + array.subs(dict(zip(array.indices, i))) for i in space_indices + ] + return dict(zip(targets, array_targets)) + + +@singledispatch +def remove_targets(expr, targets): + return (0, expr) if expr in targets else (expr, 0) + + +@remove_targets.register(sympy.Add) +def _(expr, targets): + if not any(expr.has(t) for t in targets): + return (expr, 0) + + args_b, args_F = zip(*(remove_targets(a, targets) for a in expr.args)) + return (expr.func(*args_b, evaluate=False), expr.func(*args_F, evaluate=False)) + + +@remove_targets.register(Mul) +def _(expr, targets): + if not any(expr.has(t) for t in targets): + return (expr, 0) + + args_b, args_F = zip(*[remove_targets(a, targets) if any(a.has(t) for t in targets) + else (a, a) for a in expr.args]) + return (expr.func(*args_b, evaluate=False), expr.func(*args_F, evaluate=False)) + + +@remove_targets.register(Derivative) +def _(expr, targets): + return (0, expr) if any(expr.has(t) for t in targets) else (expr, 0) + + +@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(sympy.Add) +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) + + +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 initial lowering. + NOTE: All functions used in PETSc callback functions are attached to + the `LinearSolveExpr` object, which is passed through the initial lowering + (and subsequently dropped and replaced with calls to run the solver). + Therefore, the appropriate time loop will always be correctly generated inside + the main kernel. + + Examples + -------- + >>> funcs = [ + >>> f1(t + dt, x, y), + >>> g1(t + dt, x, y), + >>> g2(t, x, y), + >>> f1(t, x, y) + >>> ] + >>> generate_time_mapper(funcs) + {t + dt: tau0, t: tau1} + + """ + time_indices = list({ + i if isinstance(d, SteppingDimension) else d + 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 dict(zip(time_indices, tau_symbs)) 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..38ac3bb9f3 --- /dev/null +++ b/devito/petsc/types/array.py @@ -0,0 +1,120 @@ +from functools import cached_property +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, as_tuple +from devito.symbolics import FieldFromComposite + + +class PETScArray(ArrayBasic, Differentiable): + """ + PETScArrays are generated by the compiler only and represent + a customised variant of ArrayBasic. + Differentiable enables compatibility with standard Function objects, + allowing for the use of the `subs` method. + + PETScArray objects represent vector objects within PETSc. + They correspond to the spatial domain of a Function-like object + provided by the user, which is passed to PETScSolve as the target. + + TODO: Potentially re-evaluate and separate into PETScFunction(Differentiable) + and then PETScArray(ArrayBasic). + """ + + _data_alignment = False + + # Default method for the finite difference approximation weights computation. + _default_fd = 'taylor' + + __rkwargs__ = (AbstractFunction.__rkwargs__ + + ('target', 'liveness', 'coefficients', 'localinfo')) + + def __init_finalize__(self, *args, **kwargs): + + self._target = kwargs.get('target') + self._ndim = kwargs['ndim'] = len(self._target.space_dimensions) + self._dimensions = kwargs['dimensions'] = self._target.space_dimensions + + super().__init_finalize__(*args, **kwargs) + + # Symbolic (finite difference) coefficients + 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._localinfo = kwargs.get('localinfo', None) + + @property + def ndim(self): + return self._ndim + + @classmethod + def __dtype_setup__(cls, **kwargs): + return kwargs['target'].dtype + + @classmethod + def __indices_setup__(cls, *args, **kwargs): + dimensions = kwargs['target'].space_dimensions + if args: + indices = args + else: + indices = dimensions + return as_tuple(dimensions), as_tuple(indices) + + def __halo_setup__(self, **kwargs): + target = kwargs['target'] + halo = [target.halo[d] for d in target.space_dimensions] + return DimensionTuple(*halo, getters=target.space_dimensions) + + @property + def dimensions(self): + return self._dimensions + + @property + def target(self): + return self._target + + @property + def coefficients(self): + """Form of the coefficients of the function.""" + return self._coefficients + + @property + def shape(self): + return self.target.grid.shape + + @property + def space_order(self): + return self.target.space_order + + @property + def localinfo(self): + return self._localinfo + + @cached_property + def _shape_with_inhalo(self): + return self.target.shape_with_inhalo + + @cached_property + def shape_allocated(self): + return self.target.shape_allocated + + @cached_property + def _C_ctype(self): + # TODO: Switch to using PetscScalar instead of float/double + # TODO: Use cat $PETSC_DIR/$PETSC_ARCH/lib/petsc/conf/petscvariables + # | grep -E "PETSC_(SCALAR|PRECISION)" to determine the precision of + # the user's PETSc configuration. + return POINTER(dtype_to_ctype(self.dtype)) + + @property + def symbolic_shape(self): + field_from_composites = [ + FieldFromComposite('g%sm' % d.name, self.localinfo) for d in self.dimensions] + # Reverse it since DMDA is setup backwards to Devito dimensions. + return DimensionTuple(*field_from_composites[::-1], getters=self.dimensions) diff --git a/devito/petsc/types/object.py b/devito/petsc/types/object.py new file mode 100644 index 0000000000..1bcfb3a6cf --- /dev/null +++ b/devito/petsc/types/object.py @@ -0,0 +1,211 @@ +from ctypes import POINTER + +from devito.tools import CustomDtype, dtype_to_cstr +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). This is the primary DM instance + created within the main kernel and linked to the SNES + solver using `SNESSetDM`. + """ + 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 _C_free(self): + return petsc_call('DMDestroy', [Byref(self.function)]) + + @property + def _C_free_priority(self): + return 3 + + +class CallbackDM(LocalObject): + """ + PETSc Data Management object (DM). This is the DM instance + accessed within the callback functions via `SNESGetDM`. + """ + dtype = CustomDtype('DM') + + def __init__(self, *args, stencil_width=None, **kwargs): + super().__init__(*args, **kwargs) + self._stencil_width = stencil_width + + @property + def stencil_width(self): + return self._stencil_width + + +class Mat(LocalObject): + """ + PETSc Matrix object (Mat). + """ + 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 local vector object (Vec). + A local vector has ghost locations that contain values that are + owned by other MPI ranks. + """ + dtype = CustomDtype('Vec') + + +class GlobalVec(LocalObject): + """ + PETSc global vector object (Vec). + A global vector is a parallel vector that has no duplicate values + between MPI ranks. A global vector has no ghost locations. + """ + dtype = CustomDtype('Vec') + + @property + def _C_free(self): + return petsc_call('VecDestroy', [Byref(self.function)]) + + @property + def _C_free_priority(self): + return 0 + + +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): + """ + A void pointer used to satisfy the function + signature of the `FormFunction` callback. + """ + 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): + """ + Fields within the struct that are updated during the time loop. + These are not set in the `PopulateMatContext` callback. + """ + return [f for f in self.fields + if isinstance(f, (ModuloDimension, TimeDimension))] + + @property + def callback_fields(self): + """ + Fields within the struct that are initialized in the `PopulateMatContext` + callback. These fields are not updated in the time loop. + """ + return [f for f in self.fields if f not in self.time_dim_fields] + + @property + def _C_ctype(self): + return POINTER(self.dtype) if self.liveness == \ + '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 new file mode 100644 index 0000000000..1a4a778c9e --- /dev/null +++ b/devito/petsc/types/types.py @@ -0,0 +1,131 @@ +import sympy + +from devito.tools import Reconstructable, sympy_mutex + + +class LinearSolveExpr(sympy.Function, Reconstructable): + """ + A symbolic expression passed through the Operator, containing the metadata + needed to execute a linear solver. Linear problems are handled with + `SNESSetType(snes, KSPONLY)`, enabling a unified interface for both + linear and nonlinear solvers. + + # TODO: extend this + defaults: + - 'ksp_type': String with the name of the PETSc Krylov method. + Default is 'gmres' (Generalized Minimal Residual Method). + https://petsc.org/main/manualpages/KSP/KSPType/ + + - 'pc_type': String with the name of the PETSc preconditioner. + Default is 'jacobi' (i.e diagonal scaling preconditioning). + https://petsc.org/main/manualpages/PC/PCType/ + + KSP tolerances: + https://petsc.org/release/manualpages/KSP/KSPSetTolerances/ + + - 'ksp_rtol': Relative convergence tolerance. Default + is 1e-5. + - 'ksp_atol': Absolute convergence for tolerance. Default + is 1e-50. + - 'ksp_divtol': Divergence tolerance, amount residual norm can + increase before `KSPConvergedDefault()` concludes + that the method is diverging. Default is 1e5. + - 'ksp_max_it': Maximum number of iterations to use. Default + is 1e4. + """ + + __rargs__ = ('expr',) + __rkwargs__ = ('target', 'solver_parameters', 'matvecs', + 'formfuncs', 'formrhs', 'arrays', 'time_mapper', + 'localinfo') + + defaults = { + 'ksp_type': 'gmres', + 'pc_type': 'jacobi', + 'ksp_rtol': 1e-5, # Relative tolerance + 'ksp_atol': 1e-50, # Absolute tolerance + 'ksp_divtol': 1e5, # Divergence tolerance + 'ksp_max_it': 1e4 # Maximum iterations + } + + def __new__(cls, expr, target=None, solver_parameters=None, + matvecs=None, formfuncs=None, formrhs=None, + arrays=None, time_mapper=None, localinfo=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_mapper = time_mapper + obj._localinfo = localinfo + 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 + + @property + def time_mapper(self): + return self._time_mapper + + @property + def localinfo(self): + return self._localinfo + + @classmethod + def eval(cls, *args): + return None + + func = Reconstructable._rebuild diff --git a/devito/petsc/utils.py b/devito/petsc/utils.py new file mode 100644 index 0000000000..d898db23cb --- /dev/null +++ b/devito/petsc/utils.py @@ -0,0 +1,53 @@ +import os + +from devito.tools import memoized_func + + +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) + } diff --git a/devito/symbolics/extended_sympy.py b/devito/symbolics/extended_sympy.py index 8acb5cac99..dac1d20cb7 100644 --- a/devito/symbolics/extended_sympy.py +++ b/devito/symbolics/extended_sympy.py @@ -252,6 +252,10 @@ def __str__(self): def field(self): return self.call + @property + def dtype(self): + return self.field.dtype + __repr__ = __str__ diff --git a/devito/symbolics/inspection.py b/devito/symbolics/inspection.py index 0ff7fcf6ba..1fb0629dc8 100644 --- a/devito/symbolics/inspection.py +++ b/devito/symbolics/inspection.py @@ -308,11 +308,10 @@ def sympy_dtype(expr, base=None, default=None, smin=None): return default dtypes = {base} - {None} - for i in expr.free_symbols: - try: - dtypes.add(i.dtype) - except AttributeError: - pass + for i in expr.args: + dtype = getattr(i, 'dtype', None) + if dtype: + dtypes.add(dtype) dtype = infer_dtype(dtypes) diff --git a/devito/tools/dtypes_lowering.py b/devito/tools/dtypes_lowering.py index 8f3111ccaa..ce9813ce22 100644 --- a/devito/tools/dtypes_lowering.py +++ b/devito/tools/dtypes_lowering.py @@ -18,6 +18,7 @@ 'is_external_ctype', 'infer_dtype', 'CustomDtype', 'mpi4py_mapper'] + # *** Custom np.dtypes # NOTE: the following is inspired by pyopencl.cltypes diff --git a/devito/types/equation.py b/devito/types/equation.py index b3353e2cb4..3e01130424 100644 --- a/devito/types/equation.py +++ b/devito/types/equation.py @@ -237,3 +237,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 032bca303a..1668f2c289 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): @@ -139,6 +139,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_()) @@ -149,7 +150,7 @@ def pfields(self): @property def pname(self): - return self.dtype._type_.__name__ + return self._pname @property def fields(self): @@ -232,6 +233,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.cpu b/docker/Dockerfile.cpu index 8f036badb9..04da0dcc8d 100644 --- a/docker/Dockerfile.cpu +++ b/docker/Dockerfile.cpu @@ -17,7 +17,8 @@ RUN apt-get update && \ # Install for basic base not containing it RUN apt-get install -y wget flex libnuma-dev hwloc curl cmake git \ - autoconf libtool build-essential procps software-properties-common + autoconf libtool build-essential procps software-properties-common \ + gfortran pkgconf libopenblas-serial-dev # Install gcc RUN if [ -n "$gcc" ]; then \ @@ -43,7 +44,6 @@ RUN cd /tmp && mkdir openmpi && \ cd openmpi && ./autogen.pl && \ mkdir build && cd build && \ ../configure --prefix=/opt/openmpi/ \ - --disable-mpi-fortran \ --enable-mca-no-build=btl-uct --enable-mpi1-compatibility && \ make -j ${nproc} && \ make install && \ diff --git a/docker/Dockerfile.devito b/docker/Dockerfile.devito index f167c36e54..ba41309fd9 100644 --- a/docker/Dockerfile.devito +++ b/docker/Dockerfile.devito @@ -4,8 +4,26 @@ # Base image with compilers ARG base=devitocodes/bases:cpu-gcc +ARG petscinstall="" -FROM $base AS builder +FROM $base AS copybase + +################## Install PETSc ############################################ +FROM copybase AS petsccopybase + +RUN apt-get update && apt-get install -y git && \ + python3 -m venv /venv && \ + /venv/bin/pip install --no-cache-dir --upgrade pip && \ + /venv/bin/pip install --no-cache-dir --no-binary numpy numpy && \ + mkdir -p /opt/petsc && \ + cd /opt/petsc && \ + git clone -b release https://gitlab.com/petsc/petsc.git petsc && \ + cd petsc && \ + ./configure --with-fortran-bindings=0 --with-mpi-dir=/opt/openmpi --with-openblas-include=$(pkg-config --variable=includedir openblas) --with-openblas-lib=$(pkg-config --variable=libdir openblas)/libopenblas.so PETSC_ARCH=devito_build && \ + make all + +ARG petscinstall="" +FROM ${petscinstall}copybase AS builder # User/Group Ids ARG USER_ID=1000 @@ -69,6 +87,9 @@ ARG GROUP_ID=1000 ENV HOME=/app ENV APP_HOME=/app +ENV PETSC_ARCH="devito_build" +ENV PETSC_DIR="/opt/petsc/petsc" + # Create the home directory for the new app user. # Create an app user so our program doesn't run as root. # Chown all the files to the app user. @@ -102,4 +123,3 @@ USER app EXPOSE 8888 ENTRYPOINT ["/docker-entrypoint.sh"] CMD ["/jupyter"] - diff --git a/tests/test_iet.py b/tests/test_iet.py index ff963d5518..c1c78eb7ea 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, + Definition, FindNodes) from devito.ir import SymbolRegistry from devito.passes.iet.engine import Graph from devito.passes.iet.languages.C import CDataManager @@ -128,6 +129,33 @@ def test_find_symbols_nested(mode, expected): assert [f.name for f in found] == eval(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 = 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('foo2', [ + Call('foo1', [foo0_arg]) + ]) + code1 = CGen().visit(call) + assert str(code1) == 'foo2(foo1((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'))) diff --git a/tests/test_petsc.py b/tests/test_petsc.py new file mode 100644 index 0000000000..13f0d064e5 --- /dev/null +++ b/tests/test_petsc.py @@ -0,0 +1,774 @@ +import numpy as np +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, + FindNodes, retrieve_iteration_tree) +from devito.types import Constant, CCompositeObject +from devito.passes.iet.languages.C import CDataManager +from devito.petsc.types import (DM, Mat, LocalVec, PetscMPIInt, KSP, + PC, KSPConvergedReason, PETScArray, + LinearSolveExpr) +from devito.petsc.solve import PETScSolve, separate_eqn, centre_stencil +from devito.petsc.iet.nodes import Expression + + +@skipif('petsc') +def test_petsc_local_object(): + """ + Test C++ support for PETSc LocalObjects. + """ + lo0 = DM('da', stencil_width=1) + lo1 = Mat('A') + lo2 = LocalVec('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) + + +@skipif('petsc') +def test_petsc_functions(): + """ + Test C++ support for PETScArrays. + """ + grid = Grid((2, 2)) + x, y = grid.dimensions + + f0 = Function(name='f', grid=grid, space_order=2, dtype=np.float32) + f1 = Function(name='f', grid=grid, space_order=2, dtype=np.float64) + + ptr0 = PETScArray(name='ptr0', target=f0) + ptr1 = PETScArray(name='ptr1', target=f0, is_const=True) + ptr2 = PETScArray(name='ptr2', target=f1, is_const=True) + + defn0 = Definition(ptr0) + defn1 = Definition(ptr1) + defn2 = Definition(ptr2) + + expr = DummyExpr(ptr0.indexed[x, y], ptr1.indexed[x, y] + 1) + + assert str(defn0) == 'float *restrict ptr0_vec;' + assert str(defn1) == 'const float *restrict ptr1_vec;' + assert str(defn2) == 'const double *restrict ptr2_vec;' + assert str(expr) == 'ptr0[x][y] = ptr1[x][y] + 1;' + + +@skipif('petsc') +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', target=f2) + + 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)))' + + 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' + + +@skipif('petsc') +def test_petsc_solve(): + """ + Test PETScSolve. + """ + 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, opt='noop') + + callable_roots = [meta_call.root for meta_call in op._func_table.values()] + + matvec_callback = [root for root in callable_roots if root.name == 'MyMatShellMult_0'] + + 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]) + + assert str(action_expr[-1].expr.rhs) == \ + 'x_matvec_f[x + 1, y + 2]/lctx->h_x**2' + \ + ' - 2.0*x_matvec_f[x + 2, y + 2]/lctx->h_x**2' + \ + ' + x_matvec_f[x + 3, y + 2]/lctx->h_x**2' + \ + ' + x_matvec_f[x + 2, y + 1]/lctx->h_y**2' + \ + ' - 2.0*x_matvec_f[x + 2, y + 2]/lctx->h_y**2' + \ + ' + x_matvec_f[x + 2, y + 3]/lctx->h_y**2' + + assert str(rhs_expr[-1].expr.rhs) == 'g[x + 2, y + 2]' + + # 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 + assert op.arguments().get('x_M') == 1 + + assert len(retrieve_iteration_tree(op)) == 0 + + # 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()] + + # One FormRHS, MatShellMult, FormFunction, PopulateMatContext per solve + assert len(callable_roots) == 8 + + +@skipif('petsc') +def test_petsc_cast(): + """ + Test casting of PETScArray. + """ + 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') + + cb1 = [meta_call.root for meta_call in op1._func_table.values()] + cb2 = [meta_call.root for meta_call in op2._func_table.values()] + cb3 = [meta_call.root for meta_call in op3._func_table.values()] + + assert 'float (*restrict x_matvec_f1) = ' + \ + '(float (*)) x_matvec_f1_vec;' in str(cb1[0]) + assert 'float (*restrict x_matvec_f2)[info.gxm] = ' + \ + '(float (*)[info.gxm]) x_matvec_f2_vec;' in str(cb2[0]) + assert 'float (*restrict x_matvec_f3)[info.gym][info.gxm] = ' + \ + '(float (*)[info.gym][info.gxm]) x_matvec_f3_vec;' in str(cb3[0]) + + +@skipif('petsc') +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-05, + 'ksp_atol': 1e-50, 'ksp_divtol': 100000.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_0)));' in str(op1) + + assert 'PetscCall(DMDACreate2d(PETSC_COMM_SELF,DM_BOUNDARY_GHOSTED,' + \ + 'DM_BOUNDARY_GHOSTED,DMDA_STENCIL_BOX,2,2,1,1,1,4,NULL,NULL,&(da_0)));' \ + in str(op2) + + assert 'PetscCall(DMDACreate3d(PETSC_COMM_SELF,DM_BOUNDARY_GHOSTED,' + \ + 'DM_BOUNDARY_GHOSTED,DM_BOUNDARY_GHOSTED,DMDA_STENCIL_BOX,6,5,4' + \ + ',1,1,1,1,6,NULL,NULL,NULL,&(da_0)));' in str(op3) + + +@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 'struct J_0_ctx\n{' not in ccode + assert 'struct J_0_ctx\n{' in hcode + + +@skipif('petsc') +@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. + + 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)) + + 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') +@pytest.mark.parametrize('eqn, target, expected', [ + ('Eq(f1.laplace, g1).evaluate', 'f1', + ( + 'g1(x, y)', + '-2.0*f1(x, y)/h_x**2 + f1(x - h_x, y)/h_x**2 + f1(x + h_x, y)/h_x**2 ' + '- 2.0*f1(x, y)/h_y**2 + f1(x, y - h_y)/h_y**2 + f1(x, y + h_y)/h_y**2' + )), + ('Eq(g1, f1.laplace).evaluate', 'f1', + ( + '-g1(x, y)', + '-(-2.0*f1(x, y)/h_x**2 + f1(x - h_x, y)/h_x**2 + f1(x + h_x, y)/h_x**2) ' + '- (-2.0*f1(x, y)/h_y**2 + f1(x, y - h_y)/h_y**2 + f1(x, y + h_y)/h_y**2)' + )), + ('Eq(g1, f1.laplace).evaluate', 'g1', + ( + '-2.0*f1(x, y)/h_x**2 + f1(x - h_x, y)/h_x**2 + f1(x + h_x, y)/h_x**2 ' + '- 2.0*f1(x, y)/h_y**2 + f1(x, y - h_y)/h_y**2 + f1(x, y + h_y)/h_y**2', + 'g1(x, y)' + )), + ('Eq(f1 + f1.laplace, g1).evaluate', 'f1', + ( + 'g1(x, y)', + '-2.0*f1(x, y)/h_x**2 + f1(x - h_x, y)/h_x**2 + f1(x + h_x, y)/h_x**2 - 2.0' + '*f1(x, y)/h_y**2 + f1(x, y - h_y)/h_y**2 + f1(x, y + h_y)/h_y**2 + f1(x, y)' + )), + ('Eq(g1.dx + f1.dx, g1).evaluate', 'f1', + ( + '-(-g1(x, y)/h_x + g1(x + h_x, y)/h_x) + g1(x, y)', + '-f1(x, y)/h_x + f1(x + h_x, y)/h_x' + )), + ('Eq(g1.dx + f1.dx, g1).evaluate', 'g1', + ( + '-(-f1(x, y)/h_x + f1(x + h_x, y)/h_x)', + '-g1(x, y)/h_x + g1(x + h_x, y)/h_x - g1(x, y)' + )), + ('Eq(f1 * g1.dx, g1).evaluate', 'g1', + ( + '0', '(-g1(x, y)/h_x + g1(x + h_x, y)/h_x)*f1(x, y) - g1(x, y)' + )), + ('Eq(f1 * g1.dx, g1).evaluate', 'f1', + ( + 'g1(x, y)', '(-g1(x, y)/h_x + g1(x + h_x, y)/h_x)*f1(x, y)' + )), + ('Eq((f1 * g1.dx).dy, f1).evaluate', 'f1', + ( + '0', '(-1/h_y)*(-g1(x, y)/h_x + g1(x + h_x, y)/h_x)*f1(x, y) ' + '+ (-g1(x, y + h_y)/h_x + g1(x + h_x, y + h_y)/h_x)*f1(x, y + h_y)/h_y ' + '- f1(x, y)' + )), + ('Eq((f1 * g1.dx).dy, f1).evaluate', 'g1', + ( + 'f1(x, y)', '(-1/h_y)*(-g1(x, y)/h_x + g1(x + h_x, y)/h_x)*f1(x, y) + ' + '(-g1(x, y + h_y)/h_x + g1(x + h_x, y + h_y)/h_x)*f1(x, y + h_y)/h_y' + )), + ('Eq(f2.laplace, g2).evaluate', 'g2', + ( + '-(-2.0*f2(t, x, y)/h_x**2 + f2(t, x - h_x, y)/h_x**2 + f2(t, x + h_x, y)' + '/h_x**2) - (-2.0*f2(t, x, y)/h_y**2 + f2(t, x, y - h_y)/h_y**2 + ' + 'f2(t, x, y + h_y)/h_y**2)', '-g2(t, x, y)' + )), + ('Eq(f2.laplace, g2).evaluate', 'f2', + ( + 'g2(t, x, y)', '-2.0*f2(t, x, y)/h_x**2 + f2(t, x - h_x, y)/h_x**2 + ' + 'f2(t, x + h_x, y)/h_x**2 - 2.0*f2(t, x, y)/h_y**2 + f2(t, x, y - h_y)' + '/h_y**2 + f2(t, x, y + h_y)/h_y**2' + )), + ('Eq(f2.laplace, f2).evaluate', 'f2', + ( + '0', '-2.0*f2(t, x, y)/h_x**2 + f2(t, x - h_x, y)/h_x**2 + ' + 'f2(t, x + h_x, y)/h_x**2 - 2.0*f2(t, x, y)/h_y**2 + f2(t, x, y - h_y)/h_y**2' + ' + f2(t, x, y + h_y)/h_y**2 - f2(t, x, y)' + )), + ('Eq(g2*f2.laplace, f2).evaluate', 'g2', + ( + 'f2(t, x, y)', '(-2.0*f2(t, x, y)/h_x**2 + f2(t, x - h_x, y)/h_x**2 + ' + 'f2(t, x + h_x, y)/h_x**2 - 2.0*f2(t, x, y)/h_y**2 + f2(t, x, y - h_y)/h_y**2' + ' + f2(t, x, y + h_y)/h_y**2)*g2(t, x, y)' + )), + ('Eq(f2.forward.laplace, f2).evaluate', 'f2.forward', + ( + 'f2(t, x, y)', '-2.0*f2(t + dt, x, y)/h_x**2 + f2(t + dt, x - h_x, y)/h_x**2' + ' + f2(t + dt, x + h_x, y)/h_x**2 - 2.0*f2(t + dt, x, y)/h_y**2 + ' + 'f2(t + dt, x, y - h_y)/h_y**2 + f2(t + dt, x, y + h_y)/h_y**2' + )), + ('Eq(f2.forward.laplace, f2).evaluate', 'f2', + ( + '-(-2.0*f2(t + dt, x, y)/h_x**2 + f2(t + dt, x - h_x, y)/h_x**2 + ' + 'f2(t + dt, x + h_x, y)/h_x**2) - (-2.0*f2(t + dt, x, y)/h_y**2 + ' + 'f2(t + dt, x, y - h_y)/h_y**2 + f2(t + dt, x, y + h_y)/h_y**2)', + '-f2(t, x, y)' + )), + ('Eq(f2.laplace + f2.forward.laplace, g2).evaluate', 'f2.forward', + ( + '-(-2.0*f2(t, x, y)/h_x**2 + f2(t, x - h_x, y)/h_x**2 + f2(t, x + h_x, y)/' + 'h_x**2) - (-2.0*f2(t, x, y)/h_y**2 + f2(t, x, y - h_y)/h_y**2 + ' + 'f2(t, x, y + h_y)/h_y**2) + g2(t, x, y)', '-2.0*f2(t + dt, x, y)/h_x**2 + ' + 'f2(t + dt, x - h_x, y)/h_x**2 + f2(t + dt, x + h_x, y)/h_x**2 - 2.0*' + 'f2(t + dt, x, y)/h_y**2 + f2(t + dt, x, y - h_y)/h_y**2 + ' + 'f2(t + dt, x, y + h_y)/h_y**2' + )), + ('Eq(g2.laplace, f2 + g2.forward).evaluate', 'g2.forward', + ( + '-(-2.0*g2(t, x, y)/h_x**2 + g2(t, x - h_x, y)/h_x**2 + ' + 'g2(t, x + h_x, y)/h_x**2) - (-2.0*g2(t, x, y)/h_y**2 + g2(t, x, y - h_y)' + '/h_y**2 + g2(t, x, y + h_y)/h_y**2) + f2(t, x, y)', '-g2(t + dt, x, y)' + )) +]) +def test_separate_eval_eqn(eqn, target, expected): + """ + Test the separate_eqn function on pre-evaluated equations. + This ensures that evaluated equations can be passed to PETScSolve, + allowing users to modify stencils for specific boundary conditions, + such as implementing free surface boundary conditions. + """ + grid = Grid((2, 2)) + + so = 2 + + f1 = Function(name='f1', grid=grid, space_order=so) # noqa + g1 = Function(name='g1', grid=grid, space_order=so) # noqa + + f2 = TimeFunction(name='f2', grid=grid, space_order=so) # noqa + g2 = TimeFunction(name='g2', grid=grid, space_order=so) # noqa + + b, F, _ = separate_eqn(eval(eqn), eval(target)) + expected_b, expected_F = expected + + assert str(b) == expected_b + assert str(F) == expected_F + + +@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 + + f2 = TimeFunction(name='f2', grid=grid, space_order=so) # noqa + g2 = TimeFunction(name='g2', 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_0'].root + ff = op._func_table['FormFunction_0'].root + + assert len(mv.parameters) == 3 + assert len(ff.parameters) == 4 + + assert str(mv.parameters) == '(J_0, X_global_0, Y_global_0)' + assert str(ff.parameters) == '(snes_0, X_global_0, F_global_0, dummy)' + + +@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=[2, 4, 8]) +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, mpi=True): + 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_0)));' + assert str(frees[1]) == 'PetscCall(VecDestroy(&(x_global_0)));' + assert str(frees[2]) == 'PetscCall(MatDestroy(&(J_0)));' + assert str(frees[3]) == 'PetscCall(SNESDestroy(&(snes_0)));' + assert str(frees[4]) == 'PetscCall(DMDestroy(&(da_0)));' + + +@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_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) + with switchconfig(openmp=False): + op1 = Operator(petsc1) + body1 = str(op1.body) + rhs1 = str(op1._func_table['FormRHS_0'].root.ccode) + + assert 'ctx0.t0 = t0' in body1 + assert 'ctx0.t1 = t1' not in body1 + assert 'lctx->t0' in rhs1 + assert 'lctx->t1' not in rhs1 + + # 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) + with switchconfig(openmp=False): + op2 = Operator(petsc2) + body2 = str(op2.body) + rhs2 = str(op2._func_table['FormRHS_0'].root.ccode) + + assert 'ctx0.time = time' in body2 + assert 'lctx->time' in rhs2 + + # Modulo time stepping with more than one time step + # used in one of the callback functions + eq3 = Eq(v1.laplace, u1 + u1.forward) + petsc3 = PETScSolve(eq3, v1) + with switchconfig(openmp=False): + op3 = Operator(petsc3) + body3 = str(op3.body) + rhs3 = str(op3._func_table['FormRHS_0'].root.ccode) + + assert 'ctx0.t0 = t0' in body3 + assert 'ctx0.t1 = t1' in body3 + assert 'lctx->t0' in rhs3 + assert 'lctx->t1' in rhs3 + + # Multiple petsc solves within the same time loop + v2 = Function(name='v2', grid=grid, space_order=2) + eq4 = Eq(v1.laplace, u1) + petsc4 = PETScSolve(eq4, v1) + eq5 = Eq(v2.laplace, u1) + petsc5 = PETScSolve(eq5, v2) + with switchconfig(openmp=False): + op4 = Operator(petsc4 + petsc5) + body4 = str(op4.body) + + assert 'ctx0.t0 = t0' in body4 + assert body4.count('ctx0.t0 = t0') == 1 diff --git a/tests/test_symbolics.py b/tests/test_symbolics.py index 9644270a6f..b768a4d523 100644 --- a/tests/test_symbolics.py +++ b/tests/test_symbolics.py @@ -14,10 +14,11 @@ CallFromPointer, Cast, DefFunction, FieldFromPointer, INT, FieldFromComposite, IntDiv, Namespace, Rvalue, ReservedWord, ListInitializer, uxreplace, - retrieve_derivatives, BaseCast) + retrieve_derivatives, BaseCast, sympy_dtype) from devito.tools import as_tuple from devito.types import (Array, Bundle, FIndexed, LocalObject, Object, - ComponentAccess, StencilDimension, Symbol as dSymbol) + ComponentAccess, StencilDimension, Symbol as dSymbol, + CompositeObject) from devito.types.basic import AbstractSymbol @@ -262,6 +263,17 @@ def test_field_from_pointer(): # Free symbols assert ffp1.free_symbols == {s} + # Test dtype + f = dSymbol('f') + pfields = [(f._C_name, f._C_ctype)] + struct = CompositeObject('s1', 'myStruct', pfields) + ffp4 = FieldFromPointer(f, struct) + assert str(ffp4) == 's1->f' + assert ffp4.dtype == f.dtype + expr = 1/ffp4 + dtype = sympy_dtype(expr) + assert dtype == f.dtype + def test_field_from_composite(): s = Symbol('s') @@ -306,7 +318,7 @@ def test_extended_sympy_arithmetic(): # noncommutative o = Object(name='o', dtype=c_void_p) bar = FieldFromPointer('bar', o) - assert ccode(-1 + bar) == '-1 + o->bar' + assert ccode(-1 + bar) == 'o->bar - 1' def test_integer_abs(): From 5505a122268bfdbaa782d235a0b04022c7c9b69a Mon Sep 17 00:00:00 2001 From: ZoeLeibowitz Date: Tue, 18 Feb 2025 18:51:51 +0000 Subject: [PATCH 02/84] Clean up --- .github/workflows/pytest-petsc.yml | 1 - devito/symbolics/inspection.py | 1 + tests/test_symbolics.py | 1 + 3 files changed, 2 insertions(+), 1 deletion(-) diff --git a/.github/workflows/pytest-petsc.yml b/.github/workflows/pytest-petsc.yml index 20bfc645d1..6f0a729e95 100644 --- a/.github/workflows/pytest-petsc.yml +++ b/.github/workflows/pytest-petsc.yml @@ -13,7 +13,6 @@ on: pull_request: branches: - master - - FieldFromPointer jobs: pytest: diff --git a/devito/symbolics/inspection.py b/devito/symbolics/inspection.py index 1fb0629dc8..709487ffbe 100644 --- a/devito/symbolics/inspection.py +++ b/devito/symbolics/inspection.py @@ -307,6 +307,7 @@ def sympy_dtype(expr, base=None, default=None, smin=None): if expr is None: return default + # TODO: Edit/fix/update according to PR #2513 dtypes = {base} - {None} for i in expr.args: dtype = getattr(i, 'dtype', None) diff --git a/tests/test_symbolics.py b/tests/test_symbolics.py index b768a4d523..a8d9d9b9f8 100644 --- a/tests/test_symbolics.py +++ b/tests/test_symbolics.py @@ -318,6 +318,7 @@ def test_extended_sympy_arithmetic(): # noncommutative o = Object(name='o', dtype=c_void_p) bar = FieldFromPointer('bar', o) + # TODO: Edit/fix/update according to PR #2513 assert ccode(-1 + bar) == 'o->bar - 1' From f4ee1b99f7c77708bd519c3d10796d7d442dbdc5 Mon Sep 17 00:00:00 2001 From: ZoeLeibowitz Date: Mon, 10 Mar 2025 14:05:58 +0000 Subject: [PATCH 03/84] compiler: Initialize and Finalize PETSc once and command line args --- devito/ir/equations/equation.py | 8 ++-- devito/ir/iet/algorithms.py | 4 +- devito/petsc/iet/nodes.py | 16 ++------ devito/petsc/iet/passes.py | 71 +++++++++++++++++++++------------ devito/petsc/iet/routines.py | 4 +- devito/petsc/iet/utils.py | 6 +-- devito/petsc/initialize.py | 42 +++++++++++++++++++ devito/petsc/solve.py | 4 +- devito/petsc/types/macros.py | 5 +++ devito/petsc/types/object.py | 9 ++++- devito/petsc/types/types.py | 22 +++++++++- devito/petsc/utils.py | 2 +- devito/types/equation.py | 2 +- examples/petsc/init_test.py | 8 ++++ examples/petsc/makefile | 5 +++ examples/petsc/petsc_test.py | 29 ++++++++++++++ examples/petsc/test_init.c | 26 ++++++++++++ tests/test_petsc.py | 25 ++++++++---- 18 files changed, 225 insertions(+), 63 deletions(-) create mode 100644 devito/petsc/initialize.py create mode 100644 devito/petsc/types/macros.py create mode 100644 examples/petsc/init_test.py create mode 100644 examples/petsc/makefile create mode 100644 examples/petsc/petsc_test.py create mode 100644 examples/petsc/test_init.c diff --git a/devito/ir/equations/equation.py b/devito/ir/equations/equation.py index 67ed9269a4..6b2261b51c 100644 --- a/devito/ir/equations/equation.py +++ b/devito/ir/equations/equation.py @@ -11,10 +11,10 @@ from devito.tools import Pickable, Tag, frozendict from devito.types import (Eq, Inc, ReduceMax, ReduceMin, relational_min) -from devito.types.equation import InjectSolveEq +from devito.types.equation import PetscEq __all__ = ['LoweredEq', 'ClusterizedEq', 'DummyEq', 'OpInc', 'OpMin', 'OpMax', - 'identity_mapper', 'OpInjectSolve'] + 'identity_mapper', 'OpPetsc'] class IREq(sympy.Eq, Pickable): @@ -105,7 +105,7 @@ def detect(cls, expr): Inc: OpInc, ReduceMax: OpMax, ReduceMin: OpMin, - InjectSolveEq: OpInjectSolve + PetscEq: OpPetsc } try: return reduction_mapper[type(expr)] @@ -122,7 +122,7 @@ def detect(cls, expr): OpInc = Operation('+') OpMax = Operation('max') OpMin = Operation('min') -OpInjectSolve = Operation('solve') +OpPetsc = Operation('solve') identity_mapper = { diff --git a/devito/ir/iet/algorithms.py b/devito/ir/iet/algorithms.py index 9d2db185db..52f48e28b1 100644 --- a/devito/ir/iet/algorithms.py +++ b/devito/ir/iet/algorithms.py @@ -3,7 +3,7 @@ from devito.ir.iet import (Expression, Increment, Iteration, List, Conditional, SyncSpot, Section, HaloSpot, ExpressionBundle) from devito.tools import timed_pass -from devito.petsc.types import LinearSolveExpr +from devito.petsc.types import MetaData from devito.petsc.iet.utils import petsc_iet_mapper __all__ = ['iet_build'] @@ -26,7 +26,7 @@ def iet_build(stree): for e in i.exprs: if e.is_Increment: exprs.append(Increment(e)) - elif isinstance(e.rhs, LinearSolveExpr): + elif isinstance(e.rhs, MetaData): exprs.append(petsc_iet_mapper[e.operation](e, operation=e.operation)) else: exprs.append(Expression(e, operation=e.operation)) diff --git a/devito/petsc/iet/nodes.py b/devito/petsc/iet/nodes.py index 2137495487..6157768b16 100644 --- a/devito/petsc/iet/nodes.py +++ b/devito/petsc/iet/nodes.py @@ -1,20 +1,12 @@ from devito.ir.iet import Expression, Callback, FixedArgsCallable, Call -from devito.ir.equations import OpInjectSolve +from devito.ir.equations import OpPetsc -class LinearSolverExpression(Expression): +class PetscMetaData(Expression): """ - Base class for general expressions required by a - matrix-free linear solve of the form Ax=b. + Base class for general expressions required to run a PETSc solver. """ - pass - - -class InjectSolveDummy(LinearSolverExpression): - """ - Placeholder expression to run the iterative solver. - """ - def __init__(self, expr, pragmas=None, operation=OpInjectSolve): + def __init__(self, expr, pragmas=None, operation=OpPetsc): super().__init__(expr, pragmas=pragmas, operation=operation) diff --git a/devito/petsc/iet/passes.py b/devito/petsc/iet/passes.py index 561f9fd0f4..9a48def8de 100644 --- a/devito/petsc/iet/passes.py +++ b/devito/petsc/iet/passes.py @@ -1,10 +1,15 @@ import cgen as c +import numpy as np from devito.passes.iet.engine import iet_pass -from devito.ir.iet import Transformer, MapNodes, Iteration, BlankLine +from devito.ir.iet import (Transformer, MapNodes, Iteration, BlankLine, + FindNodes, Call, CallableBody) from devito.symbolics import Byref, Macro -from devito.petsc.types import (PetscMPIInt, PetscErrorCode) -from devito.petsc.iet.nodes import InjectSolveDummy +from devito.types.basic import DataSymbol +from devito.petsc.types import (PetscMPIInt, PetscErrorCode, Initialize, + Finalize, ArgvSymbol) +from devito.petsc.types.macros import petsc_func_begin_user +from devito.petsc.iet.nodes import PetscMetaData from devito.petsc.utils import core_metadata from devito.petsc.iet.routines import (CallbackBuilder, BaseObjectBuilder, BaseSetup, Solver, TimeDependent, NonTimeDependent) @@ -14,14 +19,23 @@ @iet_pass def lower_petsc(iet, **kwargs): # Check if PETScSolve was used - injectsolve_mapper = MapNodes(Iteration, InjectSolveDummy, + injectsolve_mapper = MapNodes(Iteration, PetscMetaData, 'groupby').visit(iet) if not injectsolve_mapper: return iet, {} + metadata = core_metadata() + + data = FindNodes(PetscMetaData).visit(iet) + + if any(filter(lambda i: isinstance(i.expr.rhs, Initialize), data)): + return initialize(iet), metadata + + if any(filter(lambda i: isinstance(i.expr.rhs, Finalize), data)): + return finalize(iet), metadata + targets = [i.expr.rhs.target for (i,) in injectsolve_mapper.values()] - init = init_petsc(**kwargs) # Assumption is that all targets have the same grid so can use any target here objs = build_core_objects(targets[-1], **kwargs) @@ -47,26 +61,37 @@ def lower_petsc(iet, **kwargs): iet = Transformer(subs).visit(iet) body = core + tuple(setup) + (BlankLine,) + iet.body.body - body = iet.body._rebuild( - init=init, body=body, - frees=(petsc_call('PetscFinalize', []),) - ) + body = iet.body._rebuild(body=body) iet = iet._rebuild(body=body) - metadata = core_metadata() metadata.update({'efuncs': tuple(efuncs.values())}) return iet, metadata -def init_petsc(**kwargs): - # Initialize PETSc -> for now, assuming all solver options have to be - # specified via the parameters dict in PETScSolve - # TODO: Are users going to be able to use PETSc command line arguments? - # In firedrake, they have an options_prefix for each solver, enabling the use - # of command line options - initialize = petsc_call('PetscInitialize', [Null, Null, Null, Null]) +def initialize(iet): + # should be int because the correct type for argc is a C int + # and not a int32 + argc = DataSymbol(name='argc', dtype=np.int32) + argv = ArgvSymbol(name='argv') + Help = Macro('help') - return petsc_func_begin_user, initialize + help_string = c.Line(r'static char help[] = "This is help text.\n";') + + init_body = petsc_call('PetscInitialize', [Byref(argc), Byref(argv), Null, Help]) + init_body = CallableBody( + body=(petsc_func_begin_user, help_string, init_body), + retstmt=(Call('PetscFunctionReturn', arguments=[0]),) + ) + return iet._rebuild(body=init_body) + + +def finalize(iet): + finalize_body = petsc_call('PetscFinalize', []) + finalize_body = CallableBody( + body=(petsc_func_begin_user, finalize_body), + retstmt=(Call('PetscFunctionReturn', arguments=[0]),) + ) + return iet._rebuild(body=finalize_body) def make_core_petsc_calls(objs, **kwargs): @@ -76,10 +101,7 @@ def make_core_petsc_calls(objs, **kwargs): def build_core_objects(target, **kwargs): - if kwargs['options']['mpi']: - communicator = target.grid.distributor._obj_comm - else: - communicator = 'PETSC_COMM_SELF' + communicator = 'PETSC_COMM_WORLD' return { 'size': PetscMPIInt(name='size'), @@ -128,9 +150,6 @@ def __init__(self, injectsolve, objs, iters, **kwargs): ) +# Move these to types folder Null = Macro('NULL') void = 'void' - - -# TODO: Don't use c.Line here? -petsc_func_begin_user = c.Line('PetscFunctionBeginUser;') diff --git a/devito/petsc/iet/routines.py b/devito/petsc/iet/routines.py index 81547ac64b..18572a6cbd 100644 --- a/devito/petsc/iet/routines.py +++ b/devito/petsc/iet/routines.py @@ -14,7 +14,7 @@ from devito.petsc.types import PETScArray from devito.petsc.iet.nodes import (PETScCallable, FormFunctionCallback, - MatVecCallback, InjectSolveDummy) + MatVecCallback, PetscMetaData) from devito.petsc.iet.utils import petsc_call, petsc_struct from devito.petsc.utils import solver_mapper from devito.petsc.types import (DM, CallbackDM, Mat, LocalVec, GlobalVec, KSP, PC, @@ -768,7 +768,7 @@ def _spatial_loop_nest(self, iters, injectsolve): spatial_body = [] for tree in retrieve_iteration_tree(iters[0]): root = filter_iterations(tree, key=lambda i: i.dim.is_Space)[0] - if injectsolve in FindNodes(InjectSolveDummy).visit(root): + if injectsolve in FindNodes(PetscMetaData).visit(root): spatial_body.append(root) return spatial_body diff --git a/devito/petsc/iet/utils.py b/devito/petsc/iet/utils.py index adcf709eab..4f825ad196 100644 --- a/devito/petsc/iet/utils.py +++ b/devito/petsc/iet/utils.py @@ -1,5 +1,5 @@ -from devito.petsc.iet.nodes import InjectSolveDummy, PETScCall -from devito.ir.equations import OpInjectSolve +from devito.petsc.iet.nodes import PetscMetaData, PETScCall +from devito.ir.equations import OpPetsc def petsc_call(specific_call, call_args): @@ -19,4 +19,4 @@ def petsc_struct(name, fields, pname, liveness='lazy'): # Mapping special Eq operations to their corresponding IET Expression subclass types. # These operations correspond to subclasses of Eq utilised within PETScSolve. -petsc_iet_mapper = {OpInjectSolve: InjectSolveDummy} +petsc_iet_mapper = {OpPetsc: PetscMetaData} diff --git a/devito/petsc/initialize.py b/devito/petsc/initialize.py new file mode 100644 index 0000000000..9126414658 --- /dev/null +++ b/devito/petsc/initialize.py @@ -0,0 +1,42 @@ +import os +import sys +from ctypes import POINTER, cast, c_char +import atexit + +from devito import Operator +from devito.types import Symbol +from devito.types.equation import PetscEq +from devito.petsc.types import Initialize, Finalize + +global _petsc_initialized +_petsc_initialized = False + + +def PetscInitialize(): + global _petsc_initialized + if not _petsc_initialized: + dummy = Symbol(name='d') + # TODO: Potentially just use cgen + the compiler machinery in Devito + # to generate these "dummy_ops" instead of using the Operator class. + # This would prevent circular imports when initializing during import + # from the PETSc module. + op_init = Operator( + [PetscEq(dummy, Initialize(dummy))], + name='kernel_init', opt='noop' + ) + op_finalize = Operator( + [PetscEq(dummy, Finalize(dummy))], + name='kernel_finalize', opt='noop' + ) + + # `argv_bytes` must be a list so the memory address persists + # `os.fsencode` should be preferred over `string().encode('utf-8')` + # in case there is some system specific encoding in use + argv_bytes = list(map(os.fsencode, sys.argv)) + argv_pointer = (POINTER(c_char)*len(sys.argv))( + *map(lambda s: cast(s, POINTER(c_char)), argv_bytes) + ) + op_init.apply(argc=len(sys.argv), argv=argv_pointer) + + atexit.register(op_finalize.apply) + _petsc_initialized = True diff --git a/devito/petsc/solve.py b/devito/petsc/solve.py index 4f16ded1f3..b00247ae2e 100644 --- a/devito/petsc/solve.py +++ b/devito/petsc/solve.py @@ -5,7 +5,7 @@ from devito.finite_differences.differentiable import Mul from devito.finite_differences.derivative import Derivative from devito.types import Eq, Symbol, SteppingDimension, TimeFunction -from devito.types.equation import InjectSolveEq +from devito.types.equation import PetscEq from devito.operations.solve import eval_time_derivatives from devito.symbolics import retrieve_functions from devito.tools import as_tuple @@ -65,7 +65,7 @@ def PETScSolve(eqns, target, solver_parameters=None, **kwargs): ) # Placeholder equation for inserting calls to the solver and generating # correct time loop etc - inject_solve = InjectSolveEq(target, LinearSolveExpr( + inject_solve = PetscEq(target, LinearSolveExpr( expr=tuple(funcs), target=target, solver_parameters=solver_parameters, diff --git a/devito/petsc/types/macros.py b/devito/petsc/types/macros.py new file mode 100644 index 0000000000..4355535e64 --- /dev/null +++ b/devito/petsc/types/macros.py @@ -0,0 +1,5 @@ +import cgen as c + + +# TODO: Don't use c.Line here? +petsc_func_begin_user = c.Line('PetscFunctionBeginUser;') diff --git a/devito/petsc/types/object.py b/devito/petsc/types/object.py index 1bcfb3a6cf..ae9248d2a2 100644 --- a/devito/petsc/types/object.py +++ b/devito/petsc/types/object.py @@ -1,7 +1,8 @@ -from ctypes import POINTER +from ctypes import POINTER, c_char from devito.tools import CustomDtype, dtype_to_cstr from devito.types import LocalObject, CCompositeObject, ModuloDimension, TimeDimension +from devito.types.basic import DataSymbol from devito.symbolics import Byref from devito.petsc.iet.utils import petsc_call @@ -209,3 +210,9 @@ class StartPtr(LocalObject): def __init__(self, name, dtype): super().__init__(name=name) self.dtype = CustomDtype(dtype_to_cstr(dtype), modifier=' *') + + +class ArgvSymbol(DataSymbol): + @property + def _C_ctype(self): + return POINTER(POINTER(c_char)) diff --git a/devito/petsc/types/types.py b/devito/petsc/types/types.py index 1a4a778c9e..c3b43b84af 100644 --- a/devito/petsc/types/types.py +++ b/devito/petsc/types/types.py @@ -3,7 +3,27 @@ from devito.tools import Reconstructable, sympy_mutex -class LinearSolveExpr(sympy.Function, Reconstructable): +class MetaData(sympy.Function, Reconstructable): + def __new__(cls, expr, **kwargs): + with sympy_mutex: + obj = sympy.Function.__new__(cls, expr) + obj._expr = expr + return obj + + @property + def expr(self): + return self._expr + + +class Initialize(MetaData): + pass + + +class Finalize(MetaData): + pass + + +class LinearSolveExpr(MetaData): """ A symbolic expression passed through the Operator, containing the metadata needed to execute a linear solver. Linear problems are handled with diff --git a/devito/petsc/utils.py b/devito/petsc/utils.py index d898db23cb..3dbb2c0d1a 100644 --- a/devito/petsc/utils.py +++ b/devito/petsc/utils.py @@ -45,7 +45,7 @@ def core_metadata(): lib_dir = os.path.join(petsc_dir, f'{petsc_arch}', 'lib') return { - 'includes': ('petscksp.h', 'petscsnes.h', 'petscdmda.h'), + 'includes': ('petscsnes.h', 'petscdmda.h'), 'include_dirs': include_dirs, 'libs': ('petsc'), 'lib_dirs': lib_dir, diff --git a/devito/types/equation.py b/devito/types/equation.py index 3e01130424..5f027fef62 100644 --- a/devito/types/equation.py +++ b/devito/types/equation.py @@ -239,5 +239,5 @@ class ReduceMin(Reduction): pass -class InjectSolveEq(Eq): +class PetscEq(Eq): pass diff --git a/examples/petsc/init_test.py b/examples/petsc/init_test.py new file mode 100644 index 0000000000..462865831f --- /dev/null +++ b/examples/petsc/init_test.py @@ -0,0 +1,8 @@ +import os +from devito.petsc.initialize import PetscInitialize +from devito import configuration +configuration['compiler'] = 'custom' +os.environ['CC'] = 'mpicc' + +PetscInitialize() +print("helloworld") diff --git a/examples/petsc/makefile b/examples/petsc/makefile new file mode 100644 index 0000000000..ca5d2a9f57 --- /dev/null +++ b/examples/petsc/makefile @@ -0,0 +1,5 @@ +-include ${PETSC_DIR}/petscdir.mk +include ${PETSC_DIR}/lib/petsc/conf/variables +include ${PETSC_DIR}/lib/petsc/conf/rules + +all: test diff --git a/examples/petsc/petsc_test.py b/examples/petsc/petsc_test.py new file mode 100644 index 0000000000..69b63c99ec --- /dev/null +++ b/examples/petsc/petsc_test.py @@ -0,0 +1,29 @@ +import os +import numpy as np + +from devito import (Grid, Function, Eq, Operator, configuration) +from devito.petsc import PETScSolve +from devito.petsc.initialize import PetscInitialize +configuration['compiler'] = 'custom' +os.environ['CC'] = 'mpicc' + +PetscInitialize() + + +nx = 81 +ny = 81 + +grid = Grid(shape=(nx, ny), extent=(2., 2.), dtype=np.float64) + +u = Function(name='u', grid=grid, dtype=np.float64, space_order=2) +v = Function(name='v', grid=grid, dtype=np.float64, space_order=2) + +v.data[:] = 5.0 + +eq = Eq(v, u.laplace, subdomain=grid.interior) + +petsc = PETScSolve([eq], u) + +op = Operator(petsc) + +op.apply() diff --git a/examples/petsc/test_init.c b/examples/petsc/test_init.c new file mode 100644 index 0000000000..8d92ce176e --- /dev/null +++ b/examples/petsc/test_init.c @@ -0,0 +1,26 @@ +#include + +extern PetscErrorCode PetscInit(); +extern PetscErrorCode PetscFinal(); + +int main(int argc, char **argv) +{ + PetscInit(argc, argv); + PetscPrintf(PETSC_COMM_WORLD, "Hello World!\n"); + return PetscFinalize(); +} + +PetscErrorCode PetscInit(int argc, char **argv) +{ + static char help[] = "Magic help string\n"; + PetscFunctionBeginUser; + PetscCall(PetscInitialize(&argc, &argv, NULL, help)); + PetscFunctionReturn(0); +} + +PetscErrorCode PetscFinal() +{ + PetscFunctionBeginUser; + PetscCall(PetscFinalize()); + PetscFunctionReturn(0); +} diff --git a/tests/test_petsc.py b/tests/test_petsc.py index 13f0d064e5..dafc22211c 100644 --- a/tests/test_petsc.py +++ b/tests/test_petsc.py @@ -3,7 +3,8 @@ import pytest from conftest import skipif -from devito import Grid, Function, TimeFunction, Eq, Operator, switchconfig +from devito import (Grid, Function, TimeFunction, Eq, Operator, switchconfig, + configuration) from devito.ir.iet import (Call, ElementalFunction, Definition, DummyExpr, FindNodes, retrieve_iteration_tree) from devito.types import Constant, CCompositeObject @@ -13,6 +14,16 @@ LinearSolveExpr) from devito.petsc.solve import PETScSolve, separate_eqn, centre_stencil from devito.petsc.iet.nodes import Expression +from devito.petsc.initialize import PetscInitialize + + +@skipif('petsc') +def test_petsc_initialization(): + # TODO: Temporary workaround until PETSc is automatically + # initialized + configuration['compiler'] = 'custom' + os.environ['CC'] = 'mpicc' + PetscInitialize() @skipif('petsc') @@ -254,14 +265,14 @@ def test_dmda_create(): op2 = Operator(petsc2, opt='noop') op3 = Operator(petsc3, opt='noop') - assert 'PetscCall(DMDACreate1d(PETSC_COMM_SELF,DM_BOUNDARY_GHOSTED,' + \ + assert 'PetscCall(DMDACreate1d(PETSC_COMM_WORLD,DM_BOUNDARY_GHOSTED,' + \ '2,1,2,NULL,&(da_0)));' in str(op1) - assert 'PetscCall(DMDACreate2d(PETSC_COMM_SELF,DM_BOUNDARY_GHOSTED,' + \ + assert 'PetscCall(DMDACreate2d(PETSC_COMM_WORLD,DM_BOUNDARY_GHOSTED,' + \ 'DM_BOUNDARY_GHOSTED,DMDA_STENCIL_BOX,2,2,1,1,1,4,NULL,NULL,&(da_0)));' \ in str(op2) - assert 'PetscCall(DMDACreate3d(PETSC_COMM_SELF,DM_BOUNDARY_GHOSTED,' + \ + assert 'PetscCall(DMDACreate3d(PETSC_COMM_WORLD,DM_BOUNDARY_GHOSTED,' + \ 'DM_BOUNDARY_GHOSTED,DM_BOUNDARY_GHOSTED,DMDA_STENCIL_BOX,6,5,4' + \ ',1,1,1,1,6,NULL,NULL,NULL,&(da_0)));' in str(op3) @@ -604,8 +615,7 @@ def test_petsc_struct(): @skipif('petsc') -@pytest.mark.parallel(mode=[2, 4, 8]) -def test_apply(mode): +def test_apply(): grid = Grid(shape=(13, 13), dtype=np.float64) @@ -618,8 +628,7 @@ def test_apply(mode): petsc = PETScSolve(eqn, pn) # Build the op - with switchconfig(openmp=False, mpi=True): - op = Operator(petsc) + op = Operator(petsc) # Check the Operator runs without errors. Not verifying output for # now. Need to consolidate BC implementation From 0af998db0618fc9965e87cc04743cc848103fab9 Mon Sep 17 00:00:00 2001 From: ZoeLeibowitz <90093761+ZoeLeibowitz@users.noreply.github.com> Date: Wed, 19 Mar 2025 18:39:31 +0000 Subject: [PATCH 04/84] Coupled (#41) * compiler/dsl: Add machinery to support coupled solvers --- .github/workflows/pytest-petsc.yml | 2 +- devito/ir/equations/algorithms.py | 3 +- devito/ir/iet/nodes.py | 9 +- devito/ir/iet/visitors.py | 22 +- devito/passes/iet/definitions.py | 5 +- devito/petsc/clusters.py | 2 +- devito/petsc/iet/nodes.py | 2 +- devito/petsc/iet/passes.py | 189 +++- devito/petsc/iet/routines.py | 1338 ++++++++++++++++++++-------- devito/petsc/iet/utils.py | 5 +- devito/petsc/solve.py | 222 +++-- devito/petsc/types/object.py | 208 +++-- devito/petsc/types/types.py | 216 ++++- devito/symbolics/extended_sympy.py | 40 +- devito/tools/dtypes_lowering.py | 10 +- devito/types/array.py | 15 + devito/types/object.py | 27 +- tests/test_petsc.py | 343 +++++-- 18 files changed, 2001 insertions(+), 657 deletions(-) diff --git a/.github/workflows/pytest-petsc.yml b/.github/workflows/pytest-petsc.yml index 6f0a729e95..91483e2fc4 100644 --- a/.github/workflows/pytest-petsc.yml +++ b/.github/workflows/pytest-petsc.yml @@ -64,7 +64,7 @@ jobs: - name: Test with pytest run: | - ${{ env.RUN_CMD }} pytest --cov --cov-config=.coveragerc --cov-report=xml ${{ env.TESTS }} + ${{ env.RUN_CMD }} mpiexec -n 1 pytest --cov --cov-config=.coveragerc --cov-report=xml ${{ env.TESTS }} - name: Upload coverage to Codecov if: "!contains(matrix.name, 'docker')" diff --git a/devito/ir/equations/algorithms.py b/devito/ir/equations/algorithms.py index ce844887aa..5828879012 100644 --- a/devito/ir/equations/algorithms.py +++ b/devito/ir/equations/algorithms.py @@ -191,7 +191,8 @@ def concretize_subdims(exprs, **kwargs): """ sregistry = kwargs.get('sregistry') - mapper = {} + # Update based on changes in #2509 + mapper = kwargs.get('concretize_mapper', {}) rebuilt = {} # Rebuilt implicit dims etc which are shared between dimensions _concretize_subdims(exprs, mapper, rebuilt, sregistry) diff --git a/devito/ir/iet/nodes.py b/devito/ir/iet/nodes.py index a066dc60ca..16cd6825fb 100644 --- a/devito/ir/iet/nodes.py +++ b/devito/ir/iet/nodes.py @@ -1046,7 +1046,7 @@ class Dereference(ExprStmt, Node): The following cases are supported: * `pointer` is a PointerArray or TempFunction, and `pointee` is an Array. - * `pointer` is an ArrayObject or CCompositeObject representing a pointer + * `pointer` is an ArrayObject or LocalCompositeObject representing a pointer to a C struct, and `pointee` is a field in `pointer`. * `pointer` is a Symbol with its _C_ctype deriving from ct._Pointer, and `pointee` is a Symbol representing the dereferenced value. @@ -1070,8 +1070,7 @@ def functions(self): def expr_symbols(self): ret = [] if self.pointer.is_Symbol: - assert (isinstance(self.pointer, LocalCompositeObject) or - issubclass(self.pointer._C_ctype, ctypes._Pointer)), \ + assert issubclass(self.pointer._C_ctype, ctypes._Pointer), \ "Scalar dereference must have a pointer ctype" ret.extend([self.pointer._C_symbol, self.pointee._C_symbol]) elif self.pointer.is_PointerArray or self.pointer.is_TempFunction: @@ -1080,7 +1079,9 @@ def expr_symbols(self): for i in self.pointee.symbolic_shape[1:])) ret.extend(self.pointer.free_symbols) else: - ret.extend([self.pointer.indexed, self.pointee._C_symbol]) + assert (isinstance(self.pointer, LocalCompositeObject) or + issubclass(self.pointer._C_ctype, ctypes._Pointer)) + ret.extend([self.pointer._C_symbol, self.pointee._C_symbol]) return tuple(filter_ordered(ret)) @property diff --git a/devito/ir/iet/visitors.py b/devito/ir/iet/visitors.py index 5e3a8c0efa..34ac76cf06 100644 --- a/devito/ir/iet/visitors.py +++ b/devito/ir/iet/visitors.py @@ -25,7 +25,7 @@ c_restrict_void_p, sorted_priority) from devito.types.basic import AbstractFunction, AbstractSymbol, Basic from devito.types import (ArrayObject, CompositeObject, Dimension, Pointer, - IndexedData, DeviceMap) + IndexedData, DeviceMap, LocalCompositeObject) __all__ = ['FindApplications', 'FindNodes', 'FindSections', 'FindSymbols', @@ -201,7 +201,7 @@ def _restrict_keyword(self): def _gen_struct_decl(self, obj, masked=()): """ - Convert ctypes.Struct -> cgen.Structure. + Convert ctypes.Struct and LocalCompositeObject -> cgen.Structure. """ ctype = obj._C_ctype try: @@ -213,7 +213,16 @@ def _gen_struct_decl(self, obj, masked=()): return None except TypeError: # E.g., `ctype` is of type `dtypes_lowering.CustomDtype` - return None + if isinstance(obj, LocalCompositeObject): + # TODO: Potentially re-evaluate: Setting ctype to obj allows + # _gen_struct_decl to generate a cgen.Structure from a + # LocalCompositeObject, where obj._C_ctype is a CustomDtype. + # LocalCompositeObject has a __fields__ property, + # which allows the subsequent code in this function to function + # correctly. + ctype = obj + else: + return None try: return obj._C_typedecl @@ -718,8 +727,11 @@ def _operator_typedecls(self, o, mode='all'): for i in o._func_table.values(): if not i.local: continue - typedecls.extend([self._gen_struct_decl(j) for j in i.root.parameters - if xfilter(j)]) + typedecls.extend([ + self._gen_struct_decl(j) + for j in FindSymbols().visit(i.root) + if xfilter(j) + ]) typedecls = filter_sorted(typedecls, key=lambda i: i.tpname) return typedecls diff --git a/devito/passes/iet/definitions.py b/devito/passes/iet/definitions.py index 8f65899762..859de274b6 100644 --- a/devito/passes/iet/definitions.py +++ b/devito/passes/iet/definitions.py @@ -256,9 +256,10 @@ def _alloc_object_array_on_low_lat_mem(self, site, obj, storage): """ Allocate an Array of Objects in the low latency memory. """ + frees = getattr(obj, '_C_free', None) decl = Definition(obj) - storage.update(obj, site, allocs=decl) + storage.update(obj, site, allocs=decl, frees=frees) def _alloc_pointed_array_on_high_bw_mem(self, site, obj, storage): """ @@ -335,7 +336,7 @@ def _inject_definitions(self, iet, storage): frees.extend(as_list(cbody.frees) + flatten(v.frees)) frees = sorted(frees, key=lambda x: min( (obj._C_free_priority for obj in FindSymbols().visit(x) - if obj.is_LocalObject), default=float('inf') + if obj.is_LocalType), default=float('inf') )) # maps/unmaps diff --git a/devito/petsc/clusters.py b/devito/petsc/clusters.py index 0a6679b2dd..e035ccbefc 100644 --- a/devito/petsc/clusters.py +++ b/devito/petsc/clusters.py @@ -20,7 +20,7 @@ def petsc_lift(clusters): processed = [] for c in clusters: if isinstance(c.exprs[0].rhs, LinearSolveExpr): - ispace = c.ispace.lift(c.exprs[0].rhs.target.space_dimensions) + ispace = c.ispace.lift(c.exprs[0].rhs.fielddata.space_dimensions) processed.append(c.rebuild(ispace=ispace)) else: processed.append(c) diff --git a/devito/petsc/iet/nodes.py b/devito/petsc/iet/nodes.py index 6157768b16..abb5da3acd 100644 --- a/devito/petsc/iet/nodes.py +++ b/devito/petsc/iet/nodes.py @@ -14,7 +14,7 @@ class PETScCallable(FixedArgsCallable): pass -class MatVecCallback(Callback): +class MatShellSetOp(Callback): @property def callback_form(self): param_types_str = ', '.join([str(t) for t in self.param_types]) diff --git a/devito/petsc/iet/passes.py b/devito/petsc/iet/passes.py index 9a48def8de..710f25a611 100644 --- a/devito/petsc/iet/passes.py +++ b/devito/petsc/iet/passes.py @@ -1,18 +1,27 @@ import cgen as c import numpy as np +from functools import cached_property from devito.passes.iet.engine import iet_pass from devito.ir.iet import (Transformer, MapNodes, Iteration, BlankLine, - FindNodes, Call, CallableBody) -from devito.symbolics import Byref, Macro + DummyExpr, CallableBody, List, Call, Callable, + FindNodes) +from devito.symbolics import Byref, Macro, FieldFromPointer +from devito.types import Symbol, Scalar from devito.types.basic import DataSymbol -from devito.petsc.types import (PetscMPIInt, PetscErrorCode, Initialize, - Finalize, ArgvSymbol) +from devito.tools import frozendict +from devito.petsc.types import (PetscMPIInt, PetscErrorCode, MultipleFieldData, + PointerIS, Mat, LocalVec, GlobalVec, CallbackMat, SNES, + DummyArg, PetscInt, PointerDM, PointerMat, MatReuse, + CallbackPointerIS, CallbackPointerDM, JacobianStruct, + SubMatrixStruct, Initialize, Finalize, ArgvSymbol) from devito.petsc.types.macros import petsc_func_begin_user from devito.petsc.iet.nodes import PetscMetaData from devito.petsc.utils import core_metadata -from devito.petsc.iet.routines import (CallbackBuilder, BaseObjectBuilder, BaseSetup, - Solver, TimeDependent, NonTimeDependent) +from devito.petsc.iet.routines import (CBBuilder, CCBBuilder, BaseObjectBuilder, + CoupledObjectBuilder, BaseSetup, CoupledSetup, + Solver, CoupledSolver, TimeDependent, + NonTimeDependent) from devito.petsc.iet.utils import petsc_call, petsc_call_mpi @@ -26,7 +35,6 @@ def lower_petsc(iet, **kwargs): return iet, {} metadata = core_metadata() - data = FindNodes(PetscMetaData).visit(iet) if any(filter(lambda i: isinstance(i.expr.rhs, Initialize), data)): @@ -35,10 +43,10 @@ def lower_petsc(iet, **kwargs): if any(filter(lambda i: isinstance(i.expr.rhs, Finalize), data)): return finalize(iet), metadata - targets = [i.expr.rhs.target for (i,) in injectsolve_mapper.values()] - - # Assumption is that all targets have the same grid so can use any target here - objs = build_core_objects(targets[-1], **kwargs) + unique_grids = {i.expr.rhs.grid for (i,) in injectsolve_mapper.values()} + # Assumption is that all solves are on the same grid + if len(unique_grids) > 1: + raise ValueError("All PETScSolves must use the same Grid, but multiple found.") # Create core PETSc calls (not specific to each PETScSolve) core = make_core_petsc_calls(objs, **kwargs) @@ -54,17 +62,18 @@ def lower_petsc(iet, **kwargs): setup.extend(builder.solversetup.calls) # Transform the spatial iteration loop with the calls to execute the solver - subs.update(builder.solve.mapper) + subs.update({builder.solve.spatial_body: builder.solve.calls}) efuncs.update(builder.cbbuilder.efuncs) + populate_matrix_context(efuncs, objs) + iet = Transformer(subs).visit(iet) body = core + tuple(setup) + (BlankLine,) + iet.body.body body = iet.body._rebuild(body=body) iet = iet._rebuild(body=body) metadata.update({'efuncs': tuple(efuncs.values())}) - return iet, metadata @@ -100,56 +109,140 @@ def make_core_petsc_calls(objs, **kwargs): return call_mpi, BlankLine -def build_core_objects(target, **kwargs): - communicator = 'PETSC_COMM_WORLD' - - return { - 'size': PetscMPIInt(name='size'), - 'comm': communicator, - 'err': PetscErrorCode(name='err'), - 'grid': target.grid - } - - class Builder: """ This class is designed to support future extensions, enabling different combinations of solver types, preconditioning methods, and other functionalities as needed. - The class will be extended to accommodate different solver types by returning subclasses of the objects initialised in __init__, depending on the properties of `injectsolve`. """ def __init__(self, injectsolve, objs, iters, **kwargs): + self.injectsolve = injectsolve + self.objs = objs + self.iters = iters + self.kwargs = kwargs + self.coupled = isinstance(injectsolve.expr.rhs.fielddata, MultipleFieldData) + self.args = { + 'injectsolve': self.injectsolve, + 'objs': self.objs, + 'iters': self.iters, + **self.kwargs + } + self.args['solver_objs'] = self.objbuilder.solver_objs + self.args['timedep'] = self.timedep + self.args['cbbuilder'] = self.cbbuilder + + @cached_property + def objbuilder(self): + return ( + CoupledObjectBuilder(**self.args) + if self.coupled else + BaseObjectBuilder(**self.args) + ) - # Determine the time dependency class - time_mapper = injectsolve.expr.rhs.time_mapper - timedep = TimeDependent if time_mapper else NonTimeDependent - self.timedep = timedep(injectsolve, iters, **kwargs) + @cached_property + def timedep(self): + time_mapper = self.injectsolve.expr.rhs.time_mapper + timedep_class = TimeDependent if time_mapper else NonTimeDependent + return timedep_class(**self.args) - # Objects - self.objbuilder = BaseObjectBuilder(injectsolve, **kwargs) - self.solver_objs = self.objbuilder.solver_objs + @cached_property + def cbbuilder(self): + return CCBBuilder(**self.args) if self.coupled else CBBuilder(**self.args) - # Callbacks - self.cbbuilder = CallbackBuilder( - injectsolve, objs, self.solver_objs, timedep=self.timedep, - **kwargs - ) + @cached_property + def solversetup(self): + return CoupledSetup(**self.args) if self.coupled else BaseSetup(**self.args) - # Solver setup - self.solversetup = BaseSetup( - self.solver_objs, objs, injectsolve, self.cbbuilder - ) + @cached_property + def solve(self): + return CoupledSolver(**self.args) if self.coupled else Solver(**self.args) - # Execute the solver - self.solve = Solver( - self.solver_objs, objs, injectsolve, iters, - self.cbbuilder, timedep=self.timedep - ) + +def populate_matrix_context(efuncs, objs): + if not objs['dummyefunc'] in efuncs.values(): + return + + subdms_expr = DummyExpr( + FieldFromPointer(objs['Subdms']._C_symbol, objs['ljacctx']), + objs['Subdms']._C_symbol + ) + fields_expr = DummyExpr( + FieldFromPointer(objs['Fields']._C_symbol, objs['ljacctx']), + objs['Fields']._C_symbol + ) + body = CallableBody( + List(body=[subdms_expr, fields_expr]), + init=(objs['begin_user'],), + retstmt=tuple([Call('PetscFunctionReturn', arguments=[0])]) + ) + name = 'PopulateMatContext' + efuncs[name] = Callable( + name, body, objs['err'], + parameters=[objs['ljacctx'], objs['Subdms'], objs['Fields']] + ) -# Move these to types folder +# TODO: Devito MPI + PETSc testing +# if kwargs['options']['mpi'] -> communicator = grid.distributor._obj_comm +communicator = 'PETSC_COMM_WORLD' +subdms = PointerDM(name='subdms') +fields = PointerIS(name='fields') +submats = PointerMat(name='submats') +rows = PointerIS(name='rows') +cols = PointerIS(name='cols') + + +# A static dict containing shared symbols and objects that are not +# unique to each PETScSolve. +# Many of these objects are used as arguments in callback functions to make +# the C code cleaner and more modular. This is also a step toward leveraging +# Devito's `reuse_efuncs` functionality, allowing reuse of efuncs when +# they are semantically identical. +objs = frozendict({ + 'size': PetscMPIInt(name='size'), + 'comm': communicator, + 'err': PetscErrorCode(name='err'), + 'block': CallbackMat('block'), + 'submat_arr': PointerMat(name='submat_arr'), + 'subblockrows': PetscInt('subblockrows'), + 'subblockcols': PetscInt('subblockcols'), + 'rowidx': PetscInt('rowidx'), + 'colidx': PetscInt('colidx'), + 'J': Mat('J'), + 'X': GlobalVec('X'), + 'xloc': LocalVec('xloc'), + 'Y': GlobalVec('Y'), + 'yloc': LocalVec('yloc'), + 'F': GlobalVec('F'), + 'floc': LocalVec('floc'), + 'B': GlobalVec('B'), + 'nfields': PetscInt('nfields'), + 'irow': PointerIS(name='irow'), + 'icol': PointerIS(name='icol'), + 'nsubmats': Scalar('nsubmats', dtype=np.int32), + 'matreuse': MatReuse('scall'), + 'snes': SNES('snes'), + 'rows': rows, + 'cols': cols, + 'Subdms': subdms, + 'LocalSubdms': CallbackPointerDM(name='subdms'), + 'Fields': fields, + 'LocalFields': CallbackPointerIS(name='fields'), + 'Submats': submats, + 'ljacctx': JacobianStruct( + fields=[subdms, fields, submats], modifier=' *' + ), + 'subctx': SubMatrixStruct(fields=[rows, cols]), + 'Null': Macro('NULL'), + 'dummyctx': Symbol('lctx'), + 'dummyptr': DummyArg('dummy'), + 'dummyefunc': Symbol('dummyefunc'), + 'dof': PetscInt('dof'), + 'begin_user': c.Line('PetscFunctionBeginUser;'), +}) + +# Move to macros file? Null = Macro('NULL') -void = 'void' diff --git a/devito/petsc/iet/routines.py b/devito/petsc/iet/routines.py index 18572a6cbd..a70e987e71 100644 --- a/devito/petsc/iet/routines.py +++ b/devito/petsc/iet/routines.py @@ -1,50 +1,54 @@ from collections import OrderedDict - -import cgen as c +from functools import cached_property from devito.ir.iet import (Call, FindSymbols, List, Uxreplace, CallableBody, Dereference, DummyExpr, BlankLine, Callable, FindNodes, - retrieve_iteration_tree, filter_iterations) -from devito.symbolics import (Byref, FieldFromPointer, Macro, cast_mapper, - FieldFromComposite) + retrieve_iteration_tree, filter_iterations, Iteration) +from devito.symbolics import (Byref, FieldFromPointer, cast_mapper, VOIDP, + FieldFromComposite, IntDiv, Deref, Mod) from devito.symbolics.unevaluation import Mul from devito.types.basic import AbstractFunction -from devito.types import Temp, Symbol +from devito.types import Temp, Dimension from devito.tools import filter_ordered from devito.petsc.types import PETScArray from devito.petsc.iet.nodes import (PETScCallable, FormFunctionCallback, - MatVecCallback, PetscMetaData) + MatShellSetOp, PetscMetaData) from devito.petsc.iet.utils import petsc_call, petsc_struct from devito.petsc.utils import solver_mapper -from devito.petsc.types import (DM, CallbackDM, Mat, LocalVec, GlobalVec, KSP, PC, - SNES, DummyArg, PetscInt, StartPtr) +from devito.petsc.types import (DM, Mat, LocalVec, GlobalVec, KSP, PC, SNES, + PetscInt, StartPtr, PointerIS, PointerDM, VecScatter, + DMCast, JacobianStructCast, JacobianStruct, + SubMatrixStruct, CallbackDM) -class CallbackBuilder: +class CBBuilder: """ Build IET routines to generate PETSc callback functions. """ - def __init__(self, injectsolve, objs, solver_objs, - rcompile=None, sregistry=None, timedep=None, **kwargs): + def __init__(self, **kwargs): - self.rcompile = rcompile - self.sregistry = sregistry - self.timedep = timedep - self.solver_objs = solver_objs + self.rcompile = kwargs.get('rcompile', None) + self.sregistry = kwargs.get('sregistry', None) + self.concretize_mapper = kwargs.get('concretize_mapper', {}) + self.timedep = kwargs.get('timedep') + self.objs = kwargs.get('objs') + self.solver_objs = kwargs.get('solver_objs') + self.injectsolve = kwargs.get('injectsolve') self._efuncs = OrderedDict() self._struct_params = [] - self._matvec_callback = None - self._formfunc_callback = None - self._formrhs_callback = None - self._struct_callback = None + self._main_matvec_callback = None + self._main_formfunc_callback = None + self._user_struct_callback = None + # TODO: Test pickling. The mutability of these lists + # could cause issues when pickling? + self._matvecs = [] + self._formfuncs = [] + self._formrhs = [] - self._make_core(injectsolve, objs, solver_objs) - self._main_struct(solver_objs) - self._make_struct_callback(solver_objs, objs) - self._local_struct(solver_objs) + self._make_core() self._efuncs = self._uxreplace_efuncs() @property @@ -60,85 +64,104 @@ def filtered_struct_params(self): return filter_ordered(self.struct_params) @property - def matvec_callback(self): - return self._matvec_callback + def main_matvec_callback(self): + """ + This is the matvec callback associated with the whole Jacobian i.e + is set in the main kernel via + `PetscCall(MatShellSetOperation(J,MATOP_MULT,(void (*)(void))...));` + """ + return self._matvecs[0] + + @property + def main_formfunc_callback(self): + return self._formfuncs[0] + + @property + def matvecs(self): + return self._matvecs @property - def formfunc_callback(self): - return self._formfunc_callback + def formfuncs(self): + return self._formfuncs @property - def formrhs_callback(self): - return self._formrhs_callback + def formrhs(self): + return self._formrhs @property - def struct_callback(self): - return self._struct_callback + def user_struct_callback(self): + return self._user_struct_callback - def _make_core(self, injectsolve, objs, solver_objs): - self._make_matvec(injectsolve, objs, solver_objs) - self._make_formfunc(injectsolve, objs, solver_objs) - self._make_formrhs(injectsolve, objs, solver_objs) + def _make_core(self): + fielddata = self.injectsolve.expr.rhs.fielddata + self._make_matvec(fielddata, fielddata.matvecs) + self._make_formfunc(fielddata) + self._make_formrhs(fielddata) + self._make_user_struct_callback() - def _make_matvec(self, injectsolve, objs, solver_objs): + def _make_matvec(self, fielddata, matvecs, prefix='MatMult'): # Compile matvec `eqns` into an IET via recursive compilation - irs_matvec, _ = self.rcompile(injectsolve.expr.rhs.matvecs, - options={'mpi': False}, sregistry=self.sregistry) - body_matvec = self._create_matvec_body(injectsolve, - List(body=irs_matvec.uiet.body), - solver_objs, objs) - - matvec_callback = PETScCallable( - self.sregistry.make_name(prefix='MyMatShellMult_'), body_matvec, + irs_matvec, _ = self.rcompile(matvecs, + options={'mpi': False}, sregistry=self.sregistry, + concretize_mapper=self.concretize_mapper) + body_matvec = self._create_matvec_body(List(body=irs_matvec.uiet.body), + fielddata) + + objs = self.objs + cb = PETScCallable( + self.sregistry.make_name(prefix=prefix), + body_matvec, retval=objs['err'], - parameters=( - solver_objs['Jac'], solver_objs['X_global'], solver_objs['Y_global'] - ) + parameters=(objs['J'], objs['X'], objs['Y']) ) - self._matvec_callback = matvec_callback - self._efuncs[matvec_callback.name] = matvec_callback + self._matvecs.append(cb) + self._efuncs[cb.name] = cb - def _create_matvec_body(self, injectsolve, body, solver_objs, objs): - linsolve_expr = injectsolve.expr.rhs + def _create_matvec_body(self, body, fielddata): + linsolve_expr = self.injectsolve.expr.rhs + objs = self.objs + sobjs = self.solver_objs - dmda = solver_objs['callbackdm'] + dmda = sobjs['callbackdm'] + ctx = objs['dummyctx'] + xlocal = objs['xloc'] + ylocal = objs['yloc'] + y_matvec = fielddata.arrays['y'] + x_matvec = fielddata.arrays['x'] body = self.timedep.uxreplace_time(body) - fields = self._dummy_fields(body, solver_objs) + fields = self._dummy_fields(body) - y_matvec = linsolve_expr.arrays['y_matvec'] - x_matvec = linsolve_expr.arrays['x_matvec'] - - mat_get_dm = petsc_call('MatGetDM', [solver_objs['Jac'], Byref(dmda)]) + mat_get_dm = petsc_call('MatGetDM', [objs['J'], Byref(dmda)]) dm_get_app_context = petsc_call( - 'DMGetApplicationContext', [dmda, Byref(dummyctx._C_symbol)] + 'DMGetApplicationContext', [dmda, Byref(ctx._C_symbol)] ) dm_get_local_xvec = petsc_call( - 'DMGetLocalVector', [dmda, Byref(solver_objs['X_local'])] + 'DMGetLocalVector', [dmda, Byref(xlocal)] ) global_to_local_begin = petsc_call( - 'DMGlobalToLocalBegin', [dmda, solver_objs['X_global'], - 'INSERT_VALUES', solver_objs['X_local']] + 'DMGlobalToLocalBegin', [dmda, objs['X'], + insert_vals, xlocal] ) global_to_local_end = petsc_call('DMGlobalToLocalEnd', [ - dmda, solver_objs['X_global'], 'INSERT_VALUES', solver_objs['X_local'] + dmda, objs['X'], insert_vals, xlocal ]) dm_get_local_yvec = petsc_call( - 'DMGetLocalVector', [dmda, Byref(solver_objs['Y_local'])] + 'DMGetLocalVector', [dmda, Byref(ylocal)] ) vec_get_array_y = petsc_call( - 'VecGetArray', [solver_objs['Y_local'], Byref(y_matvec._C_symbol)] + 'VecGetArray', [ylocal, Byref(y_matvec._C_symbol)] ) vec_get_array_x = petsc_call( - 'VecGetArray', [solver_objs['X_local'], Byref(x_matvec._C_symbol)] + 'VecGetArray', [xlocal, Byref(x_matvec._C_symbol)] ) dm_get_local_info = petsc_call( @@ -146,27 +169,27 @@ def _create_matvec_body(self, injectsolve, body, solver_objs, objs): ) vec_restore_array_y = petsc_call( - 'VecRestoreArray', [solver_objs['Y_local'], Byref(y_matvec._C_symbol)] + 'VecRestoreArray', [ylocal, Byref(y_matvec._C_symbol)] ) vec_restore_array_x = petsc_call( - 'VecRestoreArray', [solver_objs['X_local'], Byref(x_matvec._C_symbol)] + 'VecRestoreArray', [xlocal, Byref(x_matvec._C_symbol)] ) dm_local_to_global_begin = petsc_call('DMLocalToGlobalBegin', [ - dmda, solver_objs['Y_local'], 'INSERT_VALUES', solver_objs['Y_global'] + dmda, ylocal, insert_vals, objs['Y'] ]) dm_local_to_global_end = petsc_call('DMLocalToGlobalEnd', [ - dmda, solver_objs['Y_local'], 'INSERT_VALUES', solver_objs['Y_global'] + dmda, ylocal, insert_vals, objs['Y'] ]) dm_restore_local_xvec = petsc_call( - 'DMRestoreLocalVector', [dmda, Byref(solver_objs['X_local'])] + 'DMRestoreLocalVector', [dmda, Byref(xlocal)] ) dm_restore_local_yvec = petsc_call( - 'DMRestoreLocalVector', [dmda, Byref(solver_objs['Y_local'])] + 'DMRestoreLocalVector', [dmda, Byref(ylocal)] ) # TODO: Some of the calls are placed in the `stacks` argument of the @@ -199,84 +222,88 @@ def _create_matvec_body(self, injectsolve, body, solver_objs, objs): ) # Dereference function data in struct - dereference_funcs = [Dereference(i, dummyctx) for i in + dereference_funcs = [Dereference(i, ctx) for i in fields if isinstance(i.function, AbstractFunction)] matvec_body = CallableBody( List(body=body), - init=(petsc_func_begin_user,), + init=(objs['begin_user'],), stacks=stacks+tuple(dereference_funcs), retstmt=(Call('PetscFunctionReturn', arguments=[0]),) ) # Replace non-function data with pointer to data in struct - subs = {i._C_symbol: FieldFromPointer(i._C_symbol, dummyctx) for i in fields} + subs = {i._C_symbol: FieldFromPointer(i._C_symbol, ctx) for i in fields} matvec_body = Uxreplace(subs).visit(matvec_body) self._struct_params.extend(fields) - return matvec_body - def _make_formfunc(self, injectsolve, objs, solver_objs): + def _make_formfunc(self, fielddata): + formfuncs = fielddata.formfuncs # Compile formfunc `eqns` into an IET via recursive compilation irs_formfunc, _ = self.rcompile( - injectsolve.expr.rhs.formfuncs, - options={'mpi': False}, sregistry=self.sregistry + formfuncs, options={'mpi': False}, sregistry=self.sregistry, + concretize_mapper=self.concretize_mapper ) - body_formfunc = self._create_formfunc_body(injectsolve, - List(body=irs_formfunc.uiet.body), - solver_objs, objs) - - formfunc_callback = PETScCallable( - self.sregistry.make_name(prefix='FormFunction_'), body_formfunc, + body_formfunc = self._create_formfunc_body( + List(body=irs_formfunc.uiet.body), fielddata + ) + objs = self.objs + cb = PETScCallable( + self.sregistry.make_name(prefix='FormFunction'), + body_formfunc, retval=objs['err'], - parameters=(solver_objs['snes'], solver_objs['X_global'], - solver_objs['F_global'], dummyptr) + parameters=(objs['snes'], objs['X'], objs['F'], objs['dummyptr']) ) - self._formfunc_callback = formfunc_callback - self._efuncs[formfunc_callback.name] = formfunc_callback + self._formfuncs.append(cb) + self._efuncs[cb.name] = cb - def _create_formfunc_body(self, injectsolve, body, solver_objs, objs): - linsolve_expr = injectsolve.expr.rhs + def _create_formfunc_body(self, body, fielddata): + linsolve_expr = self.injectsolve.expr.rhs + objs = self.objs + sobjs = self.solver_objs - dmda = solver_objs['callbackdm'] + dmda = sobjs['callbackdm'] + ctx = objs['dummyctx'] body = self.timedep.uxreplace_time(body) - fields = self._dummy_fields(body, solver_objs) + fields = self._dummy_fields(body) + self._struct_params.extend(fields) - f_formfunc = linsolve_expr.arrays['f_formfunc'] - x_formfunc = linsolve_expr.arrays['x_formfunc'] + f_formfunc = fielddata.arrays['f'] + x_formfunc = fielddata.arrays['x'] - snes_get_dm = petsc_call('SNESGetDM', [solver_objs['snes'], Byref(dmda)]) + dm_cast = DummyExpr(dmda, DMCast(objs['dummyptr']), init=True) dm_get_app_context = petsc_call( - 'DMGetApplicationContext', [dmda, Byref(dummyctx._C_symbol)] + 'DMGetApplicationContext', [dmda, Byref(ctx._C_symbol)] ) dm_get_local_xvec = petsc_call( - 'DMGetLocalVector', [dmda, Byref(solver_objs['X_local'])] + 'DMGetLocalVector', [dmda, Byref(objs['xloc'])] ) global_to_local_begin = petsc_call( - 'DMGlobalToLocalBegin', [dmda, solver_objs['X_global'], - 'INSERT_VALUES', solver_objs['X_local']] + 'DMGlobalToLocalBegin', [dmda, objs['X'], + insert_vals, objs['xloc']] ) global_to_local_end = petsc_call('DMGlobalToLocalEnd', [ - dmda, solver_objs['X_global'], 'INSERT_VALUES', solver_objs['X_local'] + dmda, objs['X'], insert_vals, objs['xloc'] ]) dm_get_local_yvec = petsc_call( - 'DMGetLocalVector', [dmda, Byref(solver_objs['F_local'])] + 'DMGetLocalVector', [dmda, Byref(objs['floc'])] ) vec_get_array_y = petsc_call( - 'VecGetArray', [solver_objs['F_local'], Byref(f_formfunc._C_symbol)] + 'VecGetArray', [objs['floc'], Byref(f_formfunc._C_symbol)] ) vec_get_array_x = petsc_call( - 'VecGetArray', [solver_objs['X_local'], Byref(x_formfunc._C_symbol)] + 'VecGetArray', [objs['xloc'], Byref(x_formfunc._C_symbol)] ) dm_get_local_info = petsc_call( @@ -284,27 +311,27 @@ def _create_formfunc_body(self, injectsolve, body, solver_objs, objs): ) vec_restore_array_y = petsc_call( - 'VecRestoreArray', [solver_objs['F_local'], Byref(f_formfunc._C_symbol)] + 'VecRestoreArray', [objs['floc'], Byref(f_formfunc._C_symbol)] ) vec_restore_array_x = petsc_call( - 'VecRestoreArray', [solver_objs['X_local'], Byref(x_formfunc._C_symbol)] + 'VecRestoreArray', [objs['xloc'], Byref(x_formfunc._C_symbol)] ) dm_local_to_global_begin = petsc_call('DMLocalToGlobalBegin', [ - dmda, solver_objs['F_local'], 'INSERT_VALUES', solver_objs['F_global'] + dmda, objs['floc'], insert_vals, objs['F'] ]) dm_local_to_global_end = petsc_call('DMLocalToGlobalEnd', [ - dmda, solver_objs['F_local'], 'INSERT_VALUES', solver_objs['F_global'] + dmda, objs['floc'], insert_vals, objs['F'] ]) dm_restore_local_xvec = petsc_call( - 'DMRestoreLocalVector', [dmda, Byref(solver_objs['X_local'])] + 'DMRestoreLocalVector', [dmda, Byref(objs['xloc'])] ) dm_restore_local_yvec = petsc_call( - 'DMRestoreLocalVector', [dmda, Byref(solver_objs['F_local'])] + 'DMRestoreLocalVector', [dmda, Byref(objs['floc'])] ) body = body._rebuild( @@ -318,7 +345,7 @@ def _create_formfunc_body(self, injectsolve, body, solver_objs, objs): ) stacks = ( - snes_get_dm, + dm_cast, dm_get_app_context, dm_get_local_xvec, global_to_local_begin, @@ -330,51 +357,68 @@ def _create_formfunc_body(self, injectsolve, body, solver_objs, objs): ) # Dereference function data in struct - dereference_funcs = [Dereference(i, dummyctx) for i in + dereference_funcs = [Dereference(i, ctx) for i in fields if isinstance(i.function, AbstractFunction)] formfunc_body = CallableBody( List(body=body), - init=(petsc_func_begin_user,), + init=(objs['begin_user'],), stacks=stacks+tuple(dereference_funcs), retstmt=(Call('PetscFunctionReturn', arguments=[0]),)) # Replace non-function data with pointer to data in struct - subs = {i._C_symbol: FieldFromPointer(i._C_symbol, dummyctx) for i in fields} - formfunc_body = Uxreplace(subs).visit(formfunc_body) + subs = {i._C_symbol: FieldFromPointer(i._C_symbol, ctx) for i in fields} - self._struct_params.extend(fields) + return Uxreplace(subs).visit(formfunc_body) - return formfunc_body + def _make_formrhs(self, fielddata): + formrhs = fielddata.formrhs + sobjs = self.solver_objs - def _make_formrhs(self, injectsolve, objs, solver_objs): # Compile formrhs `eqns` into an IET via recursive compilation - irs_formrhs, _ = self.rcompile(injectsolve.expr.rhs.formrhs, - options={'mpi': False}, sregistry=self.sregistry) - body_formrhs = self._create_form_rhs_body(injectsolve, - List(body=irs_formrhs.uiet.body), - solver_objs, objs) - - formrhs_callback = PETScCallable( - self.sregistry.make_name(prefix='FormRHS_'), body_formrhs, retval=objs['err'], - parameters=( - solver_objs['snes'], solver_objs['b_local'] - ) + irs_formrhs, _ = self.rcompile( + formrhs, options={'mpi': False}, sregistry=self.sregistry, + concretize_mapper=self.concretize_mapper ) - self._formrhs_callback = formrhs_callback - self._efuncs[formrhs_callback.name] = formrhs_callback + body_formrhs = self._create_form_rhs_body( + List(body=irs_formrhs.uiet.body), fielddata + ) + objs = self.objs + cb = PETScCallable( + self.sregistry.make_name(prefix='FormRHS'), + body_formrhs, + retval=objs['err'], + parameters=(sobjs['callbackdm'], objs['B']) + ) + self._formrhs.append(cb) + self._efuncs[cb.name] = cb + + def _create_form_rhs_body(self, body, fielddata): + linsolve_expr = self.injectsolve.expr.rhs + objs = self.objs + sobjs = self.solver_objs - def _create_form_rhs_body(self, injectsolve, body, solver_objs, objs): - linsolve_expr = injectsolve.expr.rhs + dmda = sobjs['callbackdm'] + ctx = objs['dummyctx'] - dmda = solver_objs['callbackdm'] + dm_get_local = petsc_call( + 'DMGetLocalVector', [dmda, Byref(sobjs['blocal'])] + ) + + dm_global_to_local_begin = petsc_call( + 'DMGlobalToLocalBegin', [dmda, objs['B'], + insert_vals, sobjs['blocal']] + ) - snes_get_dm = petsc_call('SNESGetDM', [solver_objs['snes'], Byref(dmda)]) + dm_global_to_local_end = petsc_call('DMGlobalToLocalEnd', [ + dmda, objs['B'], insert_vals, + sobjs['blocal'] + ]) - b_arr = linsolve_expr.arrays['b_tmp'] + b_arr = fielddata.arrays['b'] vec_get_array = petsc_call( - 'VecGetArray', [solver_objs['b_local'], Byref(b_arr._C_symbol)] + 'VecGetArray', [sobjs['blocal'], Byref(b_arr._C_symbol)] ) dm_get_local_info = petsc_call( @@ -383,87 +427,85 @@ def _create_form_rhs_body(self, injectsolve, body, solver_objs, objs): body = self.timedep.uxreplace_time(body) - fields = self._dummy_fields(body, solver_objs) + fields = self._dummy_fields(body) + self._struct_params.extend(fields) dm_get_app_context = petsc_call( - 'DMGetApplicationContext', [dmda, Byref(dummyctx._C_symbol)] + 'DMGetApplicationContext', [dmda, Byref(ctx._C_symbol)] ) + dm_local_to_global_begin = petsc_call('DMLocalToGlobalBegin', [ + dmda, sobjs['blocal'], insert_vals, + objs['B'] + ]) + + dm_local_to_global_end = petsc_call('DMLocalToGlobalEnd', [ + dmda, sobjs['blocal'], insert_vals, + objs['B'] + ]) + vec_restore_array = petsc_call( - 'VecRestoreArray', [solver_objs['b_local'], Byref(b_arr._C_symbol)] + 'VecRestoreArray', [sobjs['blocal'], Byref(b_arr._C_symbol)] ) - body = body._rebuild(body=body.body + (vec_restore_array,)) + body = body._rebuild(body=body.body + ( + dm_local_to_global_begin, dm_local_to_global_end, vec_restore_array + )) stacks = ( - snes_get_dm, - dm_get_app_context, + dm_get_local, + dm_global_to_local_begin, + dm_global_to_local_end, vec_get_array, + dm_get_app_context, dm_get_local_info ) # Dereference function data in struct - dereference_funcs = [Dereference(i, dummyctx) for i in + dereference_funcs = [Dereference(i, ctx) for i in fields if isinstance(i.function, AbstractFunction)] formrhs_body = CallableBody( List(body=[body]), - init=(petsc_func_begin_user,), + init=(objs['begin_user'],), stacks=stacks+tuple(dereference_funcs), retstmt=(Call('PetscFunctionReturn', arguments=[0]),) ) # Replace non-function data with pointer to data in struct - subs = {i._C_symbol: FieldFromPointer(i._C_symbol, dummyctx) for + subs = {i._C_symbol: FieldFromPointer(i._C_symbol, ctx) for i in fields if not isinstance(i.function, AbstractFunction)} - formrhs_body = Uxreplace(subs).visit(formrhs_body) - self._struct_params.extend(fields) - - return formrhs_body - - def _local_struct(self, solver_objs): - """ - This is the struct used within callback functions, - usually accessed via DMGetApplicationContext. - """ - solver_objs['localctx'] = petsc_struct( - dummyctx.name, - self.filtered_struct_params, - solver_objs['Jac'].name+'_ctx', - liveness='eager' - ) + return Uxreplace(subs).visit(formrhs_body) - def _main_struct(self, solver_objs): + def _make_user_struct_callback(self): """ This is the struct initialised inside the main kernel and attached to the DM via DMSetApplicationContext. + # TODO: this could be common between all PETScSolves instead? """ - solver_objs['mainctx'] = petsc_struct( + mainctx = self.solver_objs['userctx'] = petsc_struct( self.sregistry.make_name(prefix='ctx'), self.filtered_struct_params, - solver_objs['Jac'].name+'_ctx' + self.sregistry.make_name(prefix='UserCtx'), ) - - def _make_struct_callback(self, solver_objs, objs): - mainctx = solver_objs['mainctx'] body = [ DummyExpr(FieldFromPointer(i._C_symbol, mainctx), i._C_symbol) for i in mainctx.callback_fields ] struct_callback_body = CallableBody( - List(body=body), init=(petsc_func_begin_user,), - retstmt=tuple([Call('PetscFunctionReturn', arguments=[0])]) + List(body=body), init=(self.objs['begin_user'],), + retstmt=(Call('PetscFunctionReturn', arguments=[0]),) ) - struct_callback = Callable( - self.sregistry.make_name(prefix='PopulateMatContext_'), - struct_callback_body, objs['err'], + cb = Callable( + self.sregistry.make_name(prefix='PopulateUserContext'), + struct_callback_body, self.objs['err'], parameters=[mainctx] ) - self._efuncs[struct_callback.name] = struct_callback - self._struct_callback = struct_callback + self._efuncs[cb.name] = cb + self._user_struct_callback = cb - def _dummy_fields(self, iet, solver_objs): + def _dummy_fields(self, iet): # Place all context data required by the shell routines into a struct fields = [f.function for f in FindSymbols('basics').visit(iet)] fields = [f for f in fields if not isinstance(f.function, (PETScArray, Temp))] @@ -473,84 +515,385 @@ def _dummy_fields(self, iet, solver_objs): return fields def _uxreplace_efuncs(self): + sobjs = self.solver_objs + luserctx = petsc_struct( + sobjs['userctx'].name, + self.filtered_struct_params, + sobjs['userctx'].pname, + modifier=' *' + ) mapper = {} - visitor = Uxreplace({dummyctx: self.solver_objs['localctx']}) + visitor = Uxreplace({self.objs['dummyctx']: luserctx}) for k, v in self._efuncs.items(): mapper.update({k: visitor.visit(v)}) return mapper +class CCBBuilder(CBBuilder): + def __init__(self, **kwargs): + self._submatrices_callback = None + super().__init__(**kwargs) + + @property + def submatrices_callback(self): + return self._submatrices_callback + + @property + def submatrices(self): + return self.injectsolve.expr.rhs.fielddata.submatrices + + @property + def main_matvec_callback(self): + """ + This is the matvec callback associated with the whole Jacobian i.e + is set in the main kernel via + `PetscCall(MatShellSetOperation(J,MATOP_MULT,(void (*)(void))MyMatShellMult));` + """ + return self._main_matvec_callback + + @property + def main_formfunc_callback(self): + return self._main_formfunc_callback + + def _make_core(self): + injectsolve = self.injectsolve + targets = injectsolve.expr.rhs.fielddata.targets + all_fielddata = injectsolve.expr.rhs.fielddata + + for t in targets: + data = all_fielddata.get_field_data(t) + self._make_formfunc(data) + self._make_formrhs(data) + + row_matvecs = all_fielddata.submatrices.submatrices[t] + for submat, mtvs in row_matvecs.items(): + if mtvs['matvecs']: + self._make_matvec(data, mtvs['matvecs'], prefix=f'{submat}_MatMult') + + self._make_user_struct_callback() + self._make_whole_matvec() + self._make_whole_formfunc() + self._create_submatrices() + self._efuncs['PopulateMatContext'] = self.objs['dummyefunc'] + + def _make_whole_matvec(self): + objs = self.objs + body = self._whole_matvec_body() + + cb = PETScCallable( + self.sregistry.make_name(prefix='WholeMatMult'), + List(body=body), + retval=objs['err'], + parameters=(objs['J'], objs['X'], objs['Y']) + ) + self._main_matvec_callback = cb + self._efuncs[cb.name] = cb + + def _whole_matvec_body(self): + objs = self.objs + sobjs = self.solver_objs + + jctx = objs['ljacctx'] + ctx_main = petsc_call('MatShellGetContext', [objs['J'], Byref(jctx)]) + + nonzero_submats = self.submatrices.nonzero_submatrix_keys + + calls = () + for sm in nonzero_submats: + idx = self.submatrices.submat_to_index[sm] + ctx = sobjs[f'{sm}ctx'] + X = sobjs[f'{sm}X'] + Y = sobjs[f'{sm}Y'] + rows = objs['rows'].base + cols = objs['cols'].base + sm_indexed = objs['Submats'].indexed[idx] + + calls += ( + DummyExpr(sobjs[sm], FieldFromPointer(sm_indexed, jctx)), + petsc_call('MatShellGetContext', [sobjs[sm], Byref(ctx)]), + petsc_call( + 'VecGetSubVector', + [objs['X'], Deref(FieldFromPointer(cols, ctx)), Byref(X)] + ), + petsc_call( + 'VecGetSubVector', + [objs['Y'], Deref(FieldFromPointer(rows, ctx)), Byref(Y)] + ), + petsc_call('MatMult', [sobjs[sm], X, Y]), + petsc_call( + 'VecRestoreSubVector', + [objs['X'], Deref(FieldFromPointer(cols, ctx)), Byref(X)] + ), + petsc_call( + 'VecRestoreSubVector', + [objs['Y'], Deref(FieldFromPointer(rows, ctx)), Byref(Y)] + ), + ) + return CallableBody( + List(body=(ctx_main, BlankLine) + calls), + init=(objs['begin_user'],), + retstmt=(Call('PetscFunctionReturn', arguments=[0]),) + ) + + def _make_whole_formfunc(self): + objs = self.objs + body = self._whole_formfunc_body() + + cb = PETScCallable( + self.sregistry.make_name(prefix='WholeFormFunc'), + List(body=body), + retval=objs['err'], + parameters=(objs['snes'], objs['X'], objs['F'], objs['dummyptr']) + ) + self._main_formfunc_callback = cb + self._efuncs[cb.name] = cb + + def _whole_formfunc_body(self): + objs = self.objs + sobjs = self.solver_objs + + ljacctx = objs['ljacctx'] + struct_cast = DummyExpr(ljacctx, JacobianStructCast(objs['dummyptr'])) + X = objs['X'] + F = objs['F'] + + targets = self.injectsolve.expr.rhs.fielddata.targets + + deref_subdms = Dereference(objs['LocalSubdms'], ljacctx) + deref_fields = Dereference(objs['LocalFields'], ljacctx) + + calls = () + for i, t in enumerate(targets): + field_ptr = FieldFromPointer(objs['LocalFields'].indexed[i], ljacctx) + x_name = f'Xglobal{t.name}' + f_name = f'Fglobal{t.name}' + + calls += ( + petsc_call('VecGetSubVector', [X, field_ptr, Byref(sobjs[x_name])]), + petsc_call('VecGetSubVector', [F, field_ptr, Byref(sobjs[f_name])]), + petsc_call(self.formfuncs[i].name, [objs['snes'], sobjs[x_name], + sobjs[f_name], VOIDP(objs['LocalSubdms'].indexed[i])]), + petsc_call('VecRestoreSubVector', [X, field_ptr, Byref(sobjs[x_name])]), + petsc_call('VecRestoreSubVector', [F, field_ptr, Byref(sobjs[f_name])]), + ) + return CallableBody( + List(body=calls + (BlankLine,)), + init=(objs['begin_user'],), + stacks=(struct_cast, deref_subdms, deref_fields), + retstmt=(Call('PetscFunctionReturn', arguments=[0]),) + ) + + def _create_submatrices(self): + body = self._submat_callback_body() + objs = self.objs + params = ( + objs['J'], + objs['nfields'], + objs['irow'], + objs['icol'], + objs['matreuse'], + objs['Submats'], + ) + cb = PETScCallable( + self.sregistry.make_name(prefix='MatCreateSubMatrices'), + body, + retval=objs['err'], + parameters=params + ) + self._submatrices_callback = cb + self._efuncs[cb.name] = cb + + def _submat_callback_body(self): + objs = self.objs + sobjs = self.solver_objs + + n_submats = DummyExpr( + objs['nsubmats'], Mul(objs['nfields'], objs['nfields']) + ) + + malloc_submats = petsc_call('PetscCalloc1', [objs['nsubmats'], objs['Submats']]) + + mat_get_dm = petsc_call('MatGetDM', [objs['J'], Byref(sobjs['callbackdm'])]) + + dm_get_app = petsc_call( + 'DMGetApplicationContext', [sobjs['callbackdm'], Byref(objs['dummyctx'])] + ) + + get_ctx = petsc_call('MatShellGetContext', [objs['J'], Byref(objs['ljacctx'])]) + + Null = objs['Null'] + dm_get_info = petsc_call( + 'DMDAGetInfo', [ + sobjs['callbackdm'], Null, Byref(sobjs['M']), Byref(sobjs['N']), + Null, Null, Null, Null, Byref(objs['dof']), Null, Null, Null, Null, Null + ] + ) + subblock_rows = DummyExpr(objs['subblockrows'], Mul(sobjs['M'], sobjs['N'])) + subblock_cols = DummyExpr(objs['subblockcols'], Mul(sobjs['M'], sobjs['N'])) + + ptr = DummyExpr(objs['submat_arr']._C_symbol, Deref(objs['Submats']), init=True) + + mat_create = petsc_call('MatCreate', [self.objs['comm'], Byref(objs['block'])]) + + mat_set_sizes = petsc_call( + 'MatSetSizes', [ + objs['block'], 'PETSC_DECIDE', 'PETSC_DECIDE', + objs['subblockrows'], objs['subblockcols'] + ] + ) + + mat_set_type = petsc_call('MatSetType', [objs['block'], 'MATSHELL']) + + malloc = petsc_call('PetscMalloc1', [1, Byref(objs['subctx'])]) + i = Dimension(name='i') + + row_idx = DummyExpr(objs['rowidx'], IntDiv(i, objs['dof'])) + col_idx = DummyExpr(objs['colidx'], Mod(i, objs['dof'])) + + deref_subdm = Dereference(objs['Subdms'], objs['ljacctx']) + + set_rows = DummyExpr( + FieldFromPointer(objs['rows'].base, objs['subctx']), + Byref(objs['irow'].indexed[objs['rowidx']]) + ) + set_cols = DummyExpr( + FieldFromPointer(objs['cols'].base, objs['subctx']), + Byref(objs['icol'].indexed[objs['colidx']]) + ) + dm_set_ctx = petsc_call( + 'DMSetApplicationContext', [ + objs['Subdms'].indexed[objs['rowidx']], objs['dummyctx'] + ] + ) + matset_dm = petsc_call('MatSetDM', [ + objs['block'], objs['Subdms'].indexed[objs['rowidx']] + ]) + + set_ctx = petsc_call('MatShellSetContext', [objs['block'], objs['subctx']]) + + mat_setup = petsc_call('MatSetUp', [objs['block']]) + + assign_block = DummyExpr(objs['submat_arr'].indexed[i], objs['block']) + + iter_body = ( + mat_create, + mat_set_sizes, + mat_set_type, + malloc, + row_idx, + col_idx, + set_rows, + set_cols, + dm_set_ctx, + matset_dm, + set_ctx, + mat_setup, + assign_block + ) + + upper_bound = objs['nsubmats'] - 1 + iteration = Iteration(List(body=iter_body), i, upper_bound) + + nonzero_submats = self.submatrices.nonzero_submatrix_keys + matvec_lookup = {mv.name.split('_')[0]: mv for mv in self.matvecs} + + matmult_op = [ + petsc_call( + 'MatShellSetOperation', + [ + objs['submat_arr'].indexed[self.submatrices.submat_to_index[sb]], + 'MATOP_MULT', + MatShellSetOp(matvec_lookup[sb].name, void, void), + ], + ) + for sb in nonzero_submats if sb in matvec_lookup + ] + + body = [ + n_submats, + malloc_submats, + mat_get_dm, + dm_get_app, + dm_get_info, + subblock_rows, + subblock_cols, + ptr, + BlankLine, + iteration, + ] + matmult_op + + return CallableBody( + List(body=tuple(body)), + init=(objs['begin_user'],), + stacks=(get_ctx, deref_subdm), + retstmt=(Call('PetscFunctionReturn', arguments=[0]),) + ) + + class BaseObjectBuilder: """ A base class for constructing objects needed for a PETSc solver. Designed to be extended by subclasses, which can override the `_extend_build` method to support specific use cases. """ - - def __init__(self, injectsolve, sregistry=None, **kwargs): - self.sregistry = sregistry - self.solver_objs = self._build(injectsolve) - - def _build(self, injectsolve): + def __init__(self, **kwargs): + self.injectsolve = kwargs.get('injectsolve') + self.objs = kwargs.get('objs') + self.sregistry = kwargs.get('sregistry') + self.fielddata = self.injectsolve.expr.rhs.fielddata + self.solver_objs = self._build() + + def _build(self): """ + # TODO: update docs Constructs the core dictionary of solver objects and allows subclasses to extend or modify it via `_extend_build`. - Returns: dict: A dictionary containing the following objects: - 'Jac' (Mat): A matrix representing the jacobian. - - 'x_global' (GlobalVec): The global solution vector. - - 'x_local' (LocalVec): The local solution vector. - - 'b_global': (GlobalVec) Global RHS vector `b`, where `F(x) = b`. - - 'b_local': (LocalVec) Local RHS vector `b`, where `F(x) = b`. + - 'xglobal' (GlobalVec): The global solution vector. + - 'xlocal' (LocalVec): The local solution vector. + - 'bglobal': (GlobalVec) Global RHS vector `b`, where `F(x) = b`. + - 'blocal': (LocalVec) Local RHS vector `b`, where `F(x) = b`. - 'ksp': (KSP) Krylov solver object that manages the linear solver. - 'pc': (PC) Preconditioner object. - 'snes': (SNES) Nonlinear solver object. - - 'F_global': (GlobalVec) Global residual vector `F`, where `F(x) = b`. - - 'F_local': (LocalVec) Local residual vector `F`, where `F(x) = b`. - - 'Y_global': (GlobalVector) The output vector populated by the - matrix-free `MyMatShellMult` callback function. - - 'Y_local': (LocalVector) The output vector populated by the matrix-free - `MyMatShellMult` callback function. - - 'X_global': (GlobalVec) Current guess for the solution, - required by the FormFunction callback. - - 'X_local': (LocalVec) Current guess for the solution, - required by the FormFunction callback. - 'localsize' (PetscInt): The local length of the solution vector. - - 'start_ptr' (StartPtr): A pointer to the beginning of the solution array - that will be updated at each time step. - 'dmda' (DM): The DMDA object associated with this solve, linked to the SNES object via `SNESSetDM`. - 'callbackdm' (CallbackDM): The DM object accessed within callback functions via `SNESGetDM`. """ - target = injectsolve.expr.rhs.target sreg = self.sregistry + targets = self.fielddata.targets base_dict = { - 'Jac': Mat(sreg.make_name(prefix='J_')), - 'x_global': GlobalVec(sreg.make_name(prefix='x_global_')), - 'x_local': LocalVec(sreg.make_name(prefix='x_local_'), liveness='eager'), - 'b_global': GlobalVec(sreg.make_name(prefix='b_global_')), - 'b_local': LocalVec(sreg.make_name(prefix='b_local_')), - 'ksp': KSP(sreg.make_name(prefix='ksp_')), - 'pc': PC(sreg.make_name(prefix='pc_')), - 'snes': SNES(sreg.make_name(prefix='snes_')), - 'F_global': GlobalVec(sreg.make_name(prefix='F_global_')), - 'F_local': LocalVec(sreg.make_name(prefix='F_local_'), liveness='eager'), - 'Y_global': GlobalVec(sreg.make_name(prefix='Y_global_')), - 'Y_local': LocalVec(sreg.make_name(prefix='Y_local_'), liveness='eager'), - 'X_global': GlobalVec(sreg.make_name(prefix='X_global_')), - 'X_local': LocalVec(sreg.make_name(prefix='X_local_'), liveness='eager'), - 'localsize': PetscInt(sreg.make_name(prefix='localsize_')), - 'start_ptr': StartPtr(sreg.make_name(prefix='start_ptr_'), target.dtype), - 'dmda': DM(sreg.make_name(prefix='da_'), liveness='eager', - stencil_width=target.space_order), - 'callbackdm': CallbackDM(sreg.make_name(prefix='dm_'), - liveness='eager', stencil_width=target.space_order), + 'Jac': Mat(sreg.make_name(prefix='J')), + 'xglobal': GlobalVec(sreg.make_name(prefix='xglobal')), + 'xlocal': LocalVec(sreg.make_name(prefix='xlocal')), + 'bglobal': GlobalVec(sreg.make_name(prefix='bglobal')), + 'blocal': LocalVec(sreg.make_name(prefix='blocal')), + 'ksp': KSP(sreg.make_name(prefix='ksp')), + 'pc': PC(sreg.make_name(prefix='pc')), + 'snes': SNES(sreg.make_name(prefix='snes')), + 'localsize': PetscInt(sreg.make_name(prefix='localsize')), + 'dmda': DM(sreg.make_name(prefix='da'), dofs=len(targets)), + 'callbackdm': CallbackDM(sreg.make_name(prefix='dm')), } - return self._extend_build(base_dict, injectsolve) + self._target_dependent(base_dict) + return self._extend_build(base_dict) - def _extend_build(self, base_dict, injectsolve): + def _target_dependent(self, base_dict): + """ + '_ptr' (StartPtr): A pointer to the beginning of the solution array + that will be updated at each time step. + """ + sreg = self.sregistry + target = self.fielddata.target + base_dict[f'{target.name}_ptr'] = StartPtr( + sreg.make_name(prefix=f'{target.name}_ptr'), target.dtype + ) + + def _extend_build(self, base_dict): """ Subclasses can override this method to extend or modify the base dictionary of solver objects. @@ -558,83 +901,173 @@ def _extend_build(self, base_dict, injectsolve): return base_dict +class CoupledObjectBuilder(BaseObjectBuilder): + def _extend_build(self, base_dict): + injectsolve = self.injectsolve + sreg = self.sregistry + objs = self.objs + targets = self.fielddata.targets + + base_dict['fields'] = PointerIS( + name=sreg.make_name(prefix='fields'), nindices=len(targets) + ) + base_dict['subdms'] = PointerDM( + name=sreg.make_name(prefix='subdms'), nindices=len(targets) + ) + base_dict['nfields'] = PetscInt(sreg.make_name(prefix='nfields')) + + space_dims = len(self.fielddata.grid.dimensions) + + dim_labels = ["M", "N", "P"] + base_dict.update({ + dim_labels[i]: PetscInt(dim_labels[i]) for i in range(space_dims) + }) + + submatrices = injectsolve.expr.rhs.fielddata.submatrices + submatrix_keys = submatrices.submatrix_keys + + base_dict['jacctx'] = JacobianStruct( + name=sreg.make_name(prefix=objs['ljacctx'].name), + fields=objs['ljacctx'].fields, + ) + + for key in submatrix_keys: + base_dict[key] = Mat(name=key) + base_dict[f'{key}ctx'] = SubMatrixStruct( + name=f'{key}ctx', + fields=objs['subctx'].fields, + ) + base_dict[f'{key}X'] = LocalVec(f'{key}X') + base_dict[f'{key}Y'] = LocalVec(f'{key}Y') + base_dict[f'{key}F'] = LocalVec(f'{key}F') + + return base_dict + + def _target_dependent(self, base_dict): + sreg = self.sregistry + targets = self.fielddata.targets + for t in targets: + name = t.name + base_dict[f'{name}_ptr'] = StartPtr( + sreg.make_name(prefix=f'{name}_ptr'), t.dtype + ) + base_dict[f'xlocal{name}'] = LocalVec( + sreg.make_name(prefix=f'xlocal{name}'), liveness='eager' + ) + base_dict[f'Fglobal{name}'] = LocalVec( + sreg.make_name(prefix=f'Fglobal{name}'), liveness='eager' + ) + base_dict[f'Xglobal{name}'] = LocalVec( + sreg.make_name(prefix=f'Xglobal{name}') + ) + base_dict[f'xglobal{name}'] = GlobalVec( + sreg.make_name(prefix=f'xglobal{name}') + ) + base_dict[f'blocal{name}'] = LocalVec( + sreg.make_name(prefix=f'blocal{name}'), liveness='eager' + ) + base_dict[f'bglobal{name}'] = GlobalVec( + sreg.make_name(prefix=f'bglobal{name}') + ) + base_dict[f'da{name}'] = DM( + sreg.make_name(prefix=f'da{name}'), liveness='eager' + ) + base_dict[f'scatter{name}'] = VecScatter( + sreg.make_name(prefix=f'scatter{name}') + ) + + class BaseSetup: - def __init__(self, solver_objs, objs, injectsolve, cbbuilder): - self.calls = self._setup(solver_objs, objs, injectsolve, cbbuilder) + def __init__(self, **kwargs): + self.injectsolve = kwargs.get('injectsolve') + self.objs = kwargs.get('objs') + self.solver_objs = kwargs.get('solver_objs') + self.cbbuilder = kwargs.get('cbbuilder') + self.fielddata = self.injectsolve.expr.rhs.fielddata + self.calls = self._setup() - def _setup(self, solver_objs, objs, injectsolve, cbbuilder): - dmda = solver_objs['dmda'] + @property + def snes_ctx(self): + """ + The [optional] context for private data for the function evaluation routine. + https://petsc.org/main/manualpages/SNES/SNESSetFunction/ + """ + return VOIDP(self.solver_objs['dmda']) - solver_params = injectsolve.expr.rhs.solver_parameters + def _setup(self): + objs = self.objs + sobjs = self.solver_objs - snes_create = petsc_call('SNESCreate', [objs['comm'], Byref(solver_objs['snes'])]) + dmda = sobjs['dmda'] - snes_set_dm = petsc_call('SNESSetDM', [solver_objs['snes'], dmda]) + solver_params = self.injectsolve.expr.rhs.solver_parameters - create_matrix = petsc_call('DMCreateMatrix', [dmda, Byref(solver_objs['Jac'])]) + snes_create = petsc_call('SNESCreate', [objs['comm'], Byref(sobjs['snes'])]) - # NOTE: Assuming all solves are linear for now. - snes_set_type = petsc_call('SNESSetType', [solver_objs['snes'], 'SNESKSPONLY']) + snes_set_dm = petsc_call('SNESSetDM', [sobjs['snes'], dmda]) + + create_matrix = petsc_call('DMCreateMatrix', [dmda, Byref(sobjs['Jac'])]) + + # NOTE: Assuming all solves are linear for now + snes_set_type = petsc_call('SNESSetType', [sobjs['snes'], 'SNESKSPONLY']) snes_set_jac = petsc_call( - 'SNESSetJacobian', [solver_objs['snes'], solver_objs['Jac'], - solver_objs['Jac'], 'MatMFFDComputeJacobian', Null] + 'SNESSetJacobian', [sobjs['snes'], sobjs['Jac'], + sobjs['Jac'], 'MatMFFDComputeJacobian', objs['Null']] ) global_x = petsc_call('DMCreateGlobalVector', - [dmda, Byref(solver_objs['x_global'])]) + [dmda, Byref(sobjs['xglobal'])]) global_b = petsc_call('DMCreateGlobalVector', - [dmda, Byref(solver_objs['b_global'])]) - - local_b = petsc_call('DMCreateLocalVector', - [dmda, Byref(solver_objs['b_local'])]) + [dmda, Byref(sobjs['bglobal'])]) snes_get_ksp = petsc_call('SNESGetKSP', - [solver_objs['snes'], Byref(solver_objs['ksp'])]) + [sobjs['snes'], Byref(sobjs['ksp'])]) ksp_set_tols = petsc_call( - 'KSPSetTolerances', [solver_objs['ksp'], solver_params['ksp_rtol'], + 'KSPSetTolerances', [sobjs['ksp'], solver_params['ksp_rtol'], solver_params['ksp_atol'], solver_params['ksp_divtol'], solver_params['ksp_max_it']] ) ksp_set_type = petsc_call( - 'KSPSetType', [solver_objs['ksp'], solver_mapper[solver_params['ksp_type']]] + 'KSPSetType', [sobjs['ksp'], solver_mapper[solver_params['ksp_type']]] ) ksp_get_pc = petsc_call( - 'KSPGetPC', [solver_objs['ksp'], Byref(solver_objs['pc'])] + 'KSPGetPC', [sobjs['ksp'], Byref(sobjs['pc'])] ) # Even though the default will be jacobi, set to PCNONE for now - pc_set_type = petsc_call('PCSetType', [solver_objs['pc'], 'PCNONE']) + pc_set_type = petsc_call('PCSetType', [sobjs['pc'], 'PCNONE']) - ksp_set_from_ops = petsc_call('KSPSetFromOptions', [solver_objs['ksp']]) + ksp_set_from_ops = petsc_call('KSPSetFromOptions', [sobjs['ksp']]) + matvec = self.cbbuilder.main_matvec_callback matvec_operation = petsc_call( 'MatShellSetOperation', - [solver_objs['Jac'], 'MATOP_MULT', - MatVecCallback(cbbuilder.matvec_callback.name, void, void)] + [sobjs['Jac'], 'MATOP_MULT', MatShellSetOp(matvec.name, void, void)] ) - + formfunc = self.cbbuilder.main_formfunc_callback formfunc_operation = petsc_call( 'SNESSetFunction', - [solver_objs['snes'], Null, - FormFunctionCallback(cbbuilder.formfunc_callback.name, void, void), Null] + [sobjs['snes'], objs['Null'], FormFunctionCallback(formfunc.name, void, void), + self.snes_ctx] ) - dmda_calls = self._create_dmda_calls(dmda, objs) + dmda_calls = self._create_dmda_calls(dmda) - mainctx = solver_objs['mainctx'] + mainctx = sobjs['userctx'] call_struct_callback = petsc_call( - cbbuilder.struct_callback.name, [Byref(mainctx)] + self.cbbuilder.user_struct_callback.name, [Byref(mainctx)] ) - calls_set_app_ctx = [ - petsc_call('DMSetApplicationContext', [dmda, Byref(mainctx)]) - ] - calls = [call_struct_callback] + calls_set_app_ctx + [BlankLine] + + # TODO: maybe don't need to explictly set this + mat_set_dm = petsc_call('MatSetDM', [sobjs['Jac'], dmda]) + + calls_set_app_ctx = petsc_call('DMSetApplicationContext', [dmda, Byref(mainctx)]) base_setup = dmda_calls + ( snes_create, @@ -644,7 +1077,6 @@ def _setup(self, solver_objs, objs, injectsolve, cbbuilder): snes_set_type, global_x, global_b, - local_b, snes_get_ksp, ksp_set_tols, ksp_set_type, @@ -653,26 +1085,29 @@ def _setup(self, solver_objs, objs, injectsolve, cbbuilder): ksp_set_from_ops, matvec_operation, formfunc_operation, - ) + tuple(calls) - - extended_setup = self._extend_setup(solver_objs, objs, injectsolve, cbbuilder) - return base_setup + tuple(extended_setup) + call_struct_callback, + mat_set_dm, + calls_set_app_ctx, + BlankLine + ) + extended_setup = self._extend_setup() + return base_setup + extended_setup - def _extend_setup(self, solver_objs, objs, injectsolve, cbbuilder): + def _extend_setup(self): """ Hook for subclasses to add additional setup calls. """ - return [] + return () - def _create_dmda_calls(self, dmda, objs): - dmda_create = self._create_dmda(dmda, objs) + def _create_dmda_calls(self, dmda): + dmda_create = self._create_dmda(dmda) dm_setup = petsc_call('DMSetUp', [dmda]) dm_mat_type = petsc_call('DMSetMatType', [dmda, 'MATSHELL']) return dmda_create, dm_setup, dm_mat_type - def _create_dmda(self, dmda, objs): - grid = objs['grid'] - + def _create_dmda(self, dmda): + objs = self.objs + grid = self.fielddata.grid nspace_dims = len(grid.dimensions) # MPI communicator @@ -692,64 +1127,128 @@ def _create_dmda(self, dmda, objs): args.extend(list(grid.distributor.topology)[::-1]) # Number of degrees of freedom per node - args.append(1) + args.append(dmda.dofs) # "Stencil width" -> size of overlap - args.append(dmda.stencil_width) - args.extend([Null]*nspace_dims) + stencil_width = self.fielddata.space_order + args.append(stencil_width) + args.extend([objs['Null']]*nspace_dims) # The distributed array object args.append(Byref(dmda)) # The PETSc call used to create the DMDA - dmda = petsc_call('DMDACreate%sd' % nspace_dims, args) + dmda = petsc_call(f'DMDACreate{nspace_dims}d', args) return dmda -class Solver: - def __init__(self, solver_objs, objs, injectsolve, iters, cbbuilder, - timedep=None, **kwargs): - self.timedep = timedep - self.calls = self._execute_solve(solver_objs, objs, injectsolve, iters, cbbuilder) - self.spatial_body = self._spatial_loop_nest(iters, injectsolve) +class CoupledSetup(BaseSetup): + @property + def snes_ctx(self): + return Byref(self.solver_objs['jacctx']) + + def _extend_setup(self): + objs = self.objs + sobjs = self.solver_objs + + dmda = sobjs['dmda'] + create_field_decomp = petsc_call( + 'DMCreateFieldDecomposition', + [dmda, Byref(sobjs['nfields']), objs['Null'], Byref(sobjs['fields']), + Byref(sobjs['subdms'])] + ) + submat_cb = self.cbbuilder.submatrices_callback + matop_create_submats_op = petsc_call( + 'MatShellSetOperation', + [sobjs['Jac'], 'MATOP_CREATE_SUBMATRICES', + MatShellSetOp(submat_cb.name, void, void)] + ) + + call_coupled_struct_callback = petsc_call( + 'PopulateMatContext', + [Byref(sobjs['jacctx']), sobjs['subdms'], sobjs['fields']] + ) + + shell_set_ctx = petsc_call( + 'MatShellSetContext', [sobjs['Jac'], Byref(sobjs['jacctx']._C_symbol)] + ) - space_iter, = self.spatial_body - self.mapper = {space_iter: self.calls} + create_submats = petsc_call( + 'MatCreateSubMatrices', + [sobjs['Jac'], sobjs['nfields'], sobjs['fields'], + sobjs['fields'], 'MAT_INITIAL_MATRIX', + Byref(FieldFromComposite(objs['Submats'].base, sobjs['jacctx']))] + ) + + targets = self.fielddata.targets + + deref_dms = [ + DummyExpr(sobjs[f'da{t.name}'], sobjs['subdms'].indexed[i]) + for i, t in enumerate(targets) + ] - def _execute_solve(self, solver_objs, objs, injectsolve, iters, cbbuilder): + xglobals = [petsc_call( + 'DMCreateGlobalVector', + [sobjs[f'da{t.name}'], Byref(sobjs[f'xglobal{t.name}'])] + ) for t in targets] + + bglobals = [petsc_call( + 'DMCreateGlobalVector', + [sobjs[f'da{t.name}'], Byref(sobjs[f'bglobal{t.name}'])] + ) for t in targets] + + return ( + create_field_decomp, + matop_create_submats_op, + call_coupled_struct_callback, + shell_set_ctx, + create_submats + ) + tuple(deref_dms) + tuple(xglobals) + tuple(bglobals) + + +class Solver: + def __init__(self, **kwargs): + self.injectsolve = kwargs.get('injectsolve') + self.objs = kwargs.get('objs') + self.solver_objs = kwargs.get('solver_objs') + self.iters = kwargs.get('iters') + self.cbbuilder = kwargs.get('cbbuilder') + self.timedep = kwargs.get('timedep') + # TODO: Should/could _execute_solve be a cached_property? + self.calls = self._execute_solve() + + def _execute_solve(self): """ Assigns the required time iterators to the struct and executes the necessary calls to execute the SNES solver. """ - struct_assignment = self.timedep.assign_time_iters(solver_objs['mainctx']) + sobjs = self.solver_objs + target = self.injectsolve.expr.rhs.fielddata.target + + struct_assignment = self.timedep.assign_time_iters(sobjs['userctx']) - rhs_callback = cbbuilder.formrhs_callback + rhs_callback = self.cbbuilder.formrhs[0] - dmda = solver_objs['dmda'] + dmda = sobjs['dmda'] - rhs_call = petsc_call(rhs_callback.name, list(rhs_callback.parameters)) + rhs_call = petsc_call(rhs_callback.name, [sobjs['dmda'], sobjs['bglobal']]) local_x = petsc_call('DMCreateLocalVector', - [dmda, Byref(solver_objs['x_local'])]) + [dmda, Byref(sobjs['xlocal'])]) - vec_replace_array = self.timedep.replace_array(solver_objs) + vec_replace_array = self.timedep.replace_array(target) dm_local_to_global_x = petsc_call( - 'DMLocalToGlobal', [dmda, solver_objs['x_local'], 'INSERT_VALUES', - solver_objs['x_global']] - ) - - dm_local_to_global_b = petsc_call( - 'DMLocalToGlobal', [dmda, solver_objs['b_local'], 'INSERT_VALUES', - solver_objs['b_global']] + 'DMLocalToGlobal', [dmda, sobjs['xlocal'], insert_vals, + sobjs['xglobal']] ) snes_solve = petsc_call('SNESSolve', [ - solver_objs['snes'], solver_objs['b_global'], solver_objs['x_global']] + sobjs['snes'], sobjs['bglobal'], sobjs['xglobal']] ) dm_global_to_local_x = petsc_call('DMGlobalToLocal', [ - dmda, solver_objs['x_global'], 'INSERT_VALUES', solver_objs['x_local']] + dmda, sobjs['xglobal'], insert_vals, sobjs['xlocal']] ) run_solver_calls = (struct_assignment,) + ( @@ -757,37 +1256,120 @@ def _execute_solve(self, solver_objs, objs, injectsolve, iters, cbbuilder): local_x ) + vec_replace_array + ( dm_local_to_global_x, - dm_local_to_global_b, snes_solve, dm_global_to_local_x, BlankLine, ) return List(body=run_solver_calls) - def _spatial_loop_nest(self, iters, injectsolve): + @cached_property + def spatial_body(self): spatial_body = [] - for tree in retrieve_iteration_tree(iters[0]): + # TODO: remove the iters[0] + for tree in retrieve_iteration_tree(self.iters[0]): root = filter_iterations(tree, key=lambda i: i.dim.is_Space)[0] - if injectsolve in FindNodes(PetscMetaData).visit(root): + if self.injectsolve in FindNodes(PetscMetaData).visit(root): spatial_body.append(root) + spatial_body, = spatial_body return spatial_body -class NonTimeDependent: - def __init__(self, injectsolve, iters, **kwargs): - self.injectsolve = injectsolve - self.iters = iters - self.kwargs = kwargs - self.origin_to_moddim = self._origin_to_moddim_mapper(iters) - self.time_idx_to_symb = injectsolve.expr.rhs.time_mapper +class CoupledSolver(Solver): + def _execute_solve(self): + """ + Assigns the required time iterators to the struct and executes + the necessary calls to execute the SNES solver. + """ + sobjs = self.solver_objs + + struct_assignment = self.timedep.assign_time_iters(sobjs['userctx']) + + rhs_callbacks = self.cbbuilder.formrhs + + xglob = sobjs['xglobal'] + bglob = sobjs['bglobal'] + + targets = self.injectsolve.expr.rhs.fielddata.targets + + # TODO: optimise the ccode generated here + pre_solve = () + post_solve = () + + for i, (c, t) in enumerate(zip(rhs_callbacks, targets)): + name = t.name + dm = sobjs[f'da{name}'] + target_xloc = sobjs[f'xlocal{name}'] + target_xglob = sobjs[f'xglobal{name}'] + target_bglob = sobjs[f'bglobal{name}'] + field = sobjs['fields'].indexed[i] + s = sobjs[f'scatter{name}'] + + pre_solve += ( + petsc_call(c.name, [dm, target_bglob]), + petsc_call('DMCreateLocalVector', [dm, Byref(target_xloc)]), + self.timedep.replace_array(t), + petsc_call( + 'DMLocalToGlobal', + [dm, target_xloc, insert_vals, target_xglob] + ), + petsc_call( + 'VecScatterCreate', + [xglob, field, target_xglob, self.objs['Null'], Byref(s)] + ), + petsc_call( + 'VecScatterBegin', + [s, target_xglob, xglob, insert_vals, sreverse] + ), + petsc_call( + 'VecScatterEnd', + [s, target_xglob, xglob, insert_vals, sreverse] + ), + petsc_call( + 'VecScatterBegin', + [s, target_bglob, bglob, insert_vals, sreverse] + ), + petsc_call( + 'VecScatterEnd', + [s, target_bglob, bglob, insert_vals, sreverse] + ), + ) - @property - def is_target_time(self): - return False + post_solve += ( + petsc_call( + 'VecScatterBegin', + [s, xglob, target_xglob, insert_vals, sforward] + ), + petsc_call( + 'VecScatterEnd', + [s, xglob, target_xglob, insert_vals, sforward] + ), + petsc_call( + 'DMGlobalToLocal', + [dm, target_xglob, insert_vals, target_xloc] + ) + ) - @property - def target(self): - return self.injectsolve.expr.rhs.target + snes_solve = (petsc_call('SNESSolve', [sobjs['snes'], bglob, xglob]),) + + return List( + body=( + (struct_assignment,) + + pre_solve + + snes_solve + + post_solve + + (BlankLine,) + ) + ) + + +class NonTimeDependent: + def __init__(self, **kwargs): + self.injectsolve = kwargs.get('injectsolve') + self.iters = kwargs.get('iters') + self.sobjs = kwargs.get('solver_objs') + self.kwargs = kwargs + self.origin_to_moddim = self._origin_to_moddim_mapper(self.iters) + self.time_idx_to_symb = self.injectsolve.expr.rhs.time_mapper def _origin_to_moddim_mapper(self, iters): return {} @@ -795,7 +1377,7 @@ def _origin_to_moddim_mapper(self, iters): def uxreplace_time(self, body): return body - def replace_array(self, solver_objs): + def replace_array(self, target): """ VecReplaceArray() is a PETSc function that allows replacing the array of a `Vec` with a user provided array. @@ -806,19 +1388,19 @@ def replace_array(self, solver_objs): Examples -------- - >>> self.target + >>> target f1(x, y) - >>> call = replace_array(solver_objs) + >>> call = replace_array(target) >>> print(call) - PetscCall(VecReplaceArray(x_local_0,f1_vec->data)); + PetscCall(VecReplaceArray(xlocal0,f1_vec->data)); """ + sobjs = self.sobjs + field_from_ptr = FieldFromPointer( - self.target.function._C_field_data, self.target.function._C_symbol + target.function._C_field_data, target.function._C_symbol ) - vec_replace_array = (petsc_call( - 'VecReplaceArray', [solver_objs['x_local'], field_from_ptr] - ),) - return vec_replace_array + xlocal = sobjs.get(f'xlocal{target.name}', sobjs['xlocal']) + return (petsc_call('VecReplaceArray', [xlocal, field_from_ptr]),) def assign_time_iters(self, struct): return [] @@ -851,25 +1433,11 @@ class TimeDependent(NonTimeDependent): - Modulo dimensions are updated in the matrix context struct at each time step and can be accessed in the callback functions where needed. """ - @property - def is_target_time(self): - return any(i.is_Time for i in self.target.dimensions) - @property def time_spacing(self): - return self.target.grid.stepping_dim.spacing + return self.injectsolve.expr.rhs.grid.stepping_dim.spacing - @property - def target_time(self): - target_time = [ - i for i, d in zip(self.target.indices, self.target.dimensions) - if d.is_Time - ] - assert len(target_time) == 1 - target_time = target_time.pop() - return target_time - - @property + @cached_property def symb_to_moddim(self): """ Maps temporary `Symbol` objects created during `PETScSolve` to their @@ -881,6 +1449,18 @@ def symb_to_moddim(self): } return {symb: self.origin_to_moddim[mapper[symb]] for symb in mapper} + def is_target_time(self, target): + return any(i.is_Time for i in target.dimensions) + + def target_time(self, target): + target_time = [ + i for i, d in zip(target.indices, target.dimensions) + if d.is_Time + ] + assert len(target_time) == 1 + target_time = target_time.pop() + return target_time + def uxreplace_time(self, body): return Uxreplace(self.symb_to_moddim).visit(body) @@ -911,7 +1491,7 @@ def _origin_to_moddim_mapper(self, iters): mapper[d] = d return mapper - def replace_array(self, solver_objs): + def replace_array(self, target): """ In the case that the actual target is time-dependent e.g a `TimeFunction`, a pointer to the first element in the array that will be updated during @@ -919,51 +1499,45 @@ def replace_array(self, solver_objs): Examples -------- - >>> self.target + >>> target f1(time + dt, x, y) - >>> calls = replace_array(solver_objs) + >>> calls = replace_array(target) >>> print(List(body=calls)) - PetscCall(VecGetSize(x_local_0,&(localsize_0))); - float * start_ptr_0 = (time + 1)*localsize_0 + (float*)(f1_vec->data); - PetscCall(VecReplaceArray(x_local_0,start_ptr_0)); + PetscCall(VecGetSize(xlocal0,&(localsize0))); + float * f1_ptr0 = (time + 1)*localsize0 + (float*)(f1_vec->data); + PetscCall(VecReplaceArray(xlocal0,f1_ptr0)); - >>> self.target + >>> target f1(t + dt, x, y) - >>> calls = replace_array(solver_objs) + >>> calls = replace_array(target) >>> print(List(body=calls)) - PetscCall(VecGetSize(x_local_0,&(localsize_0))); - float * start_ptr_0 = t1*localsize_0 + (float*)(f1_vec->data); + PetscCall(VecGetSize(xlocal0,&(localsize0))); + float * f1_ptr0 = t1*localsize0 + (float*)(f1_vec->data); + PetscCall(VecReplaceArray(xlocal0,f1_ptr0)); """ - if self.is_target_time: - mapper = {self.time_spacing: 1, -self.time_spacing: -1} - target_time = self.target_time.xreplace(mapper) - - try: - target_time = self.origin_to_moddim[target_time] - except KeyError: - pass - - start_ptr = solver_objs['start_ptr'] - - vec_get_size = petsc_call( - 'VecGetSize', [solver_objs['x_local'], Byref(solver_objs['localsize'])] - ) - - field_from_ptr = FieldFromPointer( - self.target.function._C_field_data, self.target.function._C_symbol - ) + sobjs = self.sobjs - expr = DummyExpr( - start_ptr, cast_mapper[(self.target.dtype, '*')](field_from_ptr) + - Mul(target_time, solver_objs['localsize']), init=True - ) + if self.is_target_time(target): + mapper = {self.time_spacing: 1, -self.time_spacing: -1} - vec_replace_array = petsc_call( - 'VecReplaceArray', [solver_objs['x_local'], start_ptr] + target_time = self.target_time(target).xreplace(mapper) + target_time = self.origin_to_moddim.get(target_time, target_time) + + xlocal = sobjs.get(f'xlocal{target.name}', sobjs['xlocal']) + start_ptr = sobjs[f'{target.name}_ptr'] + + return ( + petsc_call('VecGetSize', [xlocal, Byref(sobjs['localsize'])]), + DummyExpr( + start_ptr, + cast_mapper[(target.dtype, '*')]( + FieldFromPointer(target._C_field_data, target._C_symbol) + ) + Mul(target_time, sobjs['localsize']), + init=True + ), + petsc_call('VecReplaceArray', [xlocal, start_ptr]) ) - return (vec_get_size, expr, vec_replace_array) - else: - return super().replace_array(solver_objs) + return super().replace_array(target) def assign_time_iters(self, struct): """ @@ -993,11 +1567,7 @@ def assign_time_iters(self, struct): return time_iter_assignments -Null = Macro('NULL') void = 'void' -dummyctx = Symbol('lctx') -dummyptr = DummyArg('dummy') - - -# TODO: Don't use c.Line here? -petsc_func_begin_user = c.Line('PetscFunctionBeginUser;') +insert_vals = 'INSERT_VALUES' +sreverse = 'SCATTER_REVERSE' +sforward = 'SCATTER_FORWARD' diff --git a/devito/petsc/iet/utils.py b/devito/petsc/iet/utils.py index 4f825ad196..99da0468ad 100644 --- a/devito/petsc/iet/utils.py +++ b/devito/petsc/iet/utils.py @@ -10,11 +10,12 @@ def petsc_call_mpi(specific_call, call_args): return PETScCall('PetscCallMPI', [PETScCall(specific_call, arguments=call_args)]) -def petsc_struct(name, fields, pname, liveness='lazy'): +def petsc_struct(name, fields, pname, liveness='lazy', modifier=None): # TODO: Fix this circular import from devito.petsc.types.object import PETScStruct return PETScStruct(name=name, pname=pname, - fields=fields, liveness=liveness) + fields=fields, liveness=liveness, + modifier=modifier) # Mapping special Eq operations to their corresponding IET Expression subclass types. diff --git a/devito/petsc/solve.py b/devito/petsc/solve.py index b00247ae2e..971fc8678b 100644 --- a/devito/petsc/solve.py +++ b/devito/petsc/solve.py @@ -8,76 +8,154 @@ from devito.types.equation import PetscEq from devito.operations.solve import eval_time_derivatives from devito.symbolics import retrieve_functions -from devito.tools import as_tuple -from devito.petsc.types import LinearSolveExpr, PETScArray, DMDALocalInfo +from devito.tools import as_tuple, filter_ordered +from devito.petsc.types import (LinearSolveExpr, PETScArray, DMDALocalInfo, + FieldData, MultipleFieldData, SubMatrices) -__all__ = ['PETScSolve'] +__all__ = ['PETScSolve', 'EssentialBC'] -def PETScSolve(eqns, target, solver_parameters=None, **kwargs): - prefixes = ['y_matvec', 'x_matvec', 'f_formfunc', 'x_formfunc', 'b_tmp'] - - localinfo = DMDALocalInfo(name='info', liveness='eager') - - arrays = { - p: PETScArray(name='%s_%s' % (p, target.name), - target=target, - liveness='eager', - localinfo=localinfo) - for p in prefixes - } - - matvecs = [] - formfuncs = [] - formrhs = [] +def PETScSolve(target_eqns, target=None, solver_parameters=None, **kwargs): + if target is not None: + return InjectSolve(solver_parameters, {target: target_eqns}).build_eq() + else: + return InjectSolveNested(solver_parameters, target_eqns).build_eq() + + +class InjectSolve: + def __init__(self, solver_parameters=None, target_eqns=None): + self.solver_params = solver_parameters + self.time_mapper = None + self.target_eqns = target_eqns + + def build_eq(self): + target, funcs, fielddata = self.linear_solve_args() + # Placeholder equation for inserting calls to the solver + linear_solve = LinearSolveExpr( + funcs, + self.solver_params, + fielddata=fielddata, + time_mapper=self.time_mapper, + localinfo=localinfo + ) + return [PetscEq(target, linear_solve)] + + def linear_solve_args(self): + target, eqns = next(iter(self.target_eqns.items())) + eqns = as_tuple(eqns) + + funcs = get_funcs(eqns) + self.time_mapper = generate_time_mapper(funcs) + arrays = self.generate_arrays(target) + + return target, tuple(funcs), self.generate_field_data(eqns, target, arrays) + + def generate_field_data(self, eqns, target, arrays): + formfuncs, formrhs = zip( + *[self.build_function_eqns(eq, target, arrays) for eq in eqns] + ) + matvecs = [self.build_matvec_eqns(eq, target, arrays) for eq in eqns] + + return FieldData( + target=target, + matvecs=matvecs, + formfuncs=formfuncs, + formrhs=formrhs, + arrays=arrays + ) + + def build_function_eqns(self, eq, target, arrays): + b, F_target, targets = separate_eqn(eq, target) + formfunc = self.make_formfunc(eq, F_target, arrays, targets) + formrhs = self.make_rhs(eq, b, arrays) - eqns = as_tuple(eqns) + return tuple(expr.subs(self.time_mapper) for expr in (formfunc, formrhs)) - for eq in eqns: + def build_matvec_eqns(self, eq, target, arrays): b, F_target, targets = separate_eqn(eq, target) - - # TODO: Current assumption is that problem is linear and user has not provided - # a jacobian. Hence, we can use F_target to form the jac-vec product - matvecs.append(Eq( - arrays['y_matvec'], - F_target.subs(targets_to_arrays(arrays['x_matvec'], targets)), - subdomain=eq.subdomain - )) - - formfuncs.append(Eq( - arrays['f_formfunc'], - F_target.subs(targets_to_arrays(arrays['x_formfunc'], targets)), - subdomain=eq.subdomain - )) - - formrhs.append(Eq( - arrays['b_tmp'], - b, - subdomain=eq.subdomain - )) - - funcs = retrieve_functions(eqns) - time_mapper = generate_time_mapper(funcs) - - matvecs, formfuncs, formrhs = ( - [eq.xreplace(time_mapper) for eq in lst] for lst in (matvecs, formfuncs, formrhs) - ) - # Placeholder equation for inserting calls to the solver and generating - # correct time loop etc - inject_solve = PetscEq(target, LinearSolveExpr( - expr=tuple(funcs), - target=target, - solver_parameters=solver_parameters, - matvecs=matvecs, - formfuncs=formfuncs, - formrhs=formrhs, - arrays=arrays, - time_mapper=time_mapper, - localinfo=localinfo - )) - - return [inject_solve] + if not F_target: + return None + matvec = self.make_matvec(eq, F_target, arrays, targets) + return matvec.subs(self.time_mapper) + + def make_matvec(self, eq, F_target, arrays, targets): + rhs = arrays['x'] if isinstance(eq, EssentialBC) else F_target.subs( + targets_to_arrays(arrays['x'], targets) + ) + return Eq(arrays['y'], rhs, subdomain=eq.subdomain) + + def make_formfunc(self, eq, F_target, arrays, targets): + rhs = 0. if isinstance(eq, EssentialBC) else F_target.subs( + targets_to_arrays(arrays['x'], targets) + ) + return Eq(arrays['f'], rhs, subdomain=eq.subdomain) + + def make_rhs(self, eq, b, arrays): + rhs = 0. if isinstance(eq, EssentialBC) else b + return Eq(arrays['b'], rhs, subdomain=eq.subdomain) + + def generate_arrays(self, target): + return { + p: PETScArray(name=f'{p}_{target.name}', + target=target, + liveness='eager', + localinfo=localinfo) + for p in prefixes + } + + +class InjectSolveNested(InjectSolve): + def linear_solve_args(self): + combined_eqns = [] + for eqns in self.target_eqns.values(): + combined_eqns.extend(eqns) + funcs = get_funcs(combined_eqns) + self.time_mapper = generate_time_mapper(funcs) + + targets = list(self.target_eqns.keys()) + jacobian = SubMatrices(targets) + + all_data = MultipleFieldData(jacobian) + + for target, eqns in self.target_eqns.items(): + eqns = as_tuple(eqns) + arrays = self.generate_arrays(target) + + self.update_jacobian(eqns, target, jacobian, arrays) + + fielddata = self.generate_field_data( + eqns, target, arrays + ) + all_data.add_field_data(fielddata) + + return target, tuple(funcs), all_data + + def update_jacobian(self, eqns, target, jacobian, arrays): + for submat, mtvs in jacobian.submatrices[target].items(): + matvecs = [ + self.build_matvec_eqns(eq, mtvs['derivative_wrt'], arrays) + for eq in eqns + ] + # Set submatrix only if there's at least one non-zero matvec + if any(m is not None for m in matvecs): + jacobian.set_submatrix(target, submat, matvecs) + + def generate_field_data(self, eqns, target, arrays): + formfuncs, formrhs = zip( + *[self.build_function_eqns(eq, target, arrays) for eq in eqns] + ) + + return FieldData( + target=target, + formfuncs=formfuncs, + formrhs=formrhs, + arrays=arrays + ) + + +class EssentialBC(Eq): + pass def separate_eqn(eqn, target): @@ -113,15 +191,12 @@ def targets_to_arrays(array, targets): """ Map each target in `targets` to a corresponding array generated from `array`, matching the spatial indices of the target. - Example: -------- >>> array vec_u(x, y) - >>> targets {u(t + dt, x + h_x, y), u(t + dt, x - h_x, y), u(t + dt, x, y)} - >>> targets_to_arrays(array, targets) {u(t + dt, x - h_x, y): vec_u(x - h_x, y), u(t + dt, x + h_x, y): vec_u(x + h_x, y), @@ -218,7 +293,6 @@ def generate_time_mapper(funcs): (and subsequently dropped and replaced with calls to run the solver). Therefore, the appropriate time loop will always be correctly generated inside the main kernel. - Examples -------- >>> funcs = [ @@ -229,7 +303,6 @@ def generate_time_mapper(funcs): >>> ] >>> generate_time_mapper(funcs) {t + dt: tau0, t: tau1} - """ time_indices = list({ i if isinstance(d, SteppingDimension) else d @@ -239,3 +312,16 @@ def generate_time_mapper(funcs): }) tau_symbs = [Symbol('tau%d' % i) for i in range(len(time_indices))] return dict(zip(time_indices, tau_symbs)) + + +def get_funcs(eqns): + funcs = [ + func + for eq in eqns + for func in retrieve_functions(eval_time_derivatives(eq.lhs - eq.rhs)) + ] + return filter_ordered(funcs) + + +localinfo = DMDALocalInfo(name='info', liveness='eager') +prefixes = ['y', 'x', 'f', 'b'] diff --git a/devito/petsc/types/object.py b/devito/petsc/types/object.py index ae9248d2a2..9acf7def46 100644 --- a/devito/petsc/types/object.py +++ b/devito/petsc/types/object.py @@ -1,13 +1,21 @@ from ctypes import POINTER, c_char - -from devito.tools import CustomDtype, dtype_to_cstr -from devito.types import LocalObject, CCompositeObject, ModuloDimension, TimeDimension +from devito.tools import CustomDtype, dtype_to_cstr, as_tuple, CustomIntType +from devito.types import (LocalObject, LocalCompositeObject, ModuloDimension, + TimeDimension, ArrayObject, CustomDimension) +from devito.symbolics import Byref, Cast from devito.types.basic import DataSymbol -from devito.symbolics import Byref - from devito.petsc.iet.utils import petsc_call +class CallbackDM(LocalObject): + """ + PETSc Data Management object (DM). This is the DM instance + accessed within the callback functions via `SNESGetDM` and + is not destroyed during callback execution. + """ + dtype = CustomDtype('DM') + + class DM(LocalObject): """ PETSc Data Management object (DM). This is the primary DM instance @@ -16,43 +24,38 @@ class DM(LocalObject): """ dtype = CustomDtype('DM') - def __init__(self, *args, stencil_width=None, **kwargs): + def __init__(self, *args, dofs=1, **kwargs): super().__init__(*args, **kwargs) - self._stencil_width = stencil_width + self._dofs = dofs @property - def stencil_width(self): - return self._stencil_width + def dofs(self): + return self._dofs @property def _C_free(self): return petsc_call('DMDestroy', [Byref(self.function)]) + # TODO: This is growing out of hand so switch to an enumeration or something? @property def _C_free_priority(self): - return 3 + return 4 -class CallbackDM(LocalObject): - """ - PETSc Data Management object (DM). This is the DM instance - accessed within the callback functions via `SNESGetDM`. - """ - dtype = CustomDtype('DM') +class DMCast(Cast): + _base_typ = 'DM' - def __init__(self, *args, stencil_width=None, **kwargs): - super().__init__(*args, **kwargs) - self._stencil_width = stencil_width - @property - def stencil_width(self): - return self._stencil_width +class CallbackMat(LocalObject): + """ + PETSc Matrix object (Mat) used within callback functions. + These instances are not destroyed during callback execution; + instead, they are managed and destroyed in the main kernel. + """ + dtype = CustomDtype('Mat') class Mat(LocalObject): - """ - PETSc Matrix object (Mat). - """ dtype = CustomDtype('Mat') @property @@ -61,7 +64,7 @@ def _C_free(self): @property def _C_free_priority(self): - return 1 + return 2 class LocalVec(LocalObject): @@ -73,21 +76,26 @@ class LocalVec(LocalObject): dtype = CustomDtype('Vec') -class GlobalVec(LocalObject): +class CallbackGlobalVec(LocalVec): + """ + PETSc global vector object (Vec). For example, used for coupled + solves inside the `WholeFormFunc` callback. + """ + + +class GlobalVec(LocalVec): """ PETSc global vector object (Vec). A global vector is a parallel vector that has no duplicate values between MPI ranks. A global vector has no ghost locations. """ - dtype = CustomDtype('Vec') - @property def _C_free(self): return petsc_call('VecDestroy', [Byref(self.function)]) @property def _C_free_priority(self): - return 0 + return 1 class PetscMPIInt(LocalObject): @@ -103,7 +111,7 @@ class PetscInt(LocalObject): PETSc datatype used to represent `int` parameters to PETSc functions. """ - dtype = CustomDtype('PetscInt') + dtype = CustomIntType('PetscInt') class KSP(LocalObject): @@ -114,19 +122,21 @@ class KSP(LocalObject): dtype = CustomDtype('KSP') -class SNES(LocalObject): +class CallbackSNES(LocalObject): """ PETSc SNES : Non-Linear Systems Solvers. """ dtype = CustomDtype('SNES') + +class SNES(CallbackSNES): @property def _C_free(self): return petsc_call('SNESDestroy', [Byref(self.function)]) @property def _C_free_priority(self): - return 2 + return 3 class PC(LocalObject): @@ -168,18 +178,25 @@ class DummyArg(LocalObject): dtype = CustomDtype('void', modifier='*') -class PETScStruct(CCompositeObject): +class MatReuse(LocalObject): + dtype = CustomDtype('MatReuse') - __rargs__ = ('name', 'pname', 'fields') - def __init__(self, name, pname, fields, liveness='lazy'): - pfields = [(i._C_name, i._C_ctype) for i in fields] - super().__init__(name, pname, pfields, liveness) - self._fields = fields +class VecScatter(LocalObject): + dtype = CustomDtype('VecScatter') - @property - def fields(self): - return self._fields + +class StartPtr(LocalObject): + def __init__(self, name, dtype): + super().__init__(name=name) + self.dtype = CustomDtype(dtype_to_cstr(dtype), modifier=' *') + + +class SingleIS(LocalObject): + dtype = CustomDtype('IS') + + +class PETScStruct(LocalCompositeObject): @property def time_dim_fields(self): @@ -198,18 +215,109 @@ def callback_fields(self): """ return [f for f in self.fields if f not in self.time_dim_fields] + _C_modifier = ' *' + + +class JacobianStruct(PETScStruct): + def __init__(self, name='jctx', pname='JacobianCtx', fields=None, + modifier='', liveness='lazy'): + super().__init__(name, pname, fields, modifier, liveness) + _C_modifier = None + + +class SubMatrixStruct(PETScStruct): + def __init__(self, name='subctx', pname='SubMatrixCtx', fields=None, + modifier=' *', liveness='lazy'): + super().__init__(name, pname, fields, modifier, liveness) + _C_modifier = None + + +class JacobianStructCast(Cast): + _base_typ = 'struct JacobianCtx *' + + +class PETScArrayObject(ArrayObject): + _data_alignment = False + + def __init_finalize__(self, *args, **kwargs): + self._nindices = kwargs.pop('nindices', 1) + super().__init_finalize__(*args, **kwargs) + + @classmethod + def __indices_setup__(cls, **kwargs): + try: + return as_tuple(kwargs['dimensions']), as_tuple(kwargs['dimensions']) + except KeyError: + nindices = kwargs.get('nindices', 1) + dim = CustomDimension(name='d', symbolic_size=nindices) + return (dim,), (dim,) + @property - def _C_ctype(self): - return POINTER(self.dtype) if self.liveness == \ - 'eager' else self.dtype + def dim(self): + assert len(self.dimensions) == 1 + return self.dimensions[0] - _C_modifier = ' *' + @property + def nindices(self): + return self._nindices + @property + def _C_name(self): + return self.name -class StartPtr(LocalObject): - def __init__(self, name, dtype): - super().__init__(name=name) - self.dtype = CustomDtype(dtype_to_cstr(dtype), modifier=' *') + @property + def _mem_stack(self): + return False + + @property + def _C_free_priority(self): + return 0 + + +class CallbackPointerIS(PETScArrayObject): + """ + Index set object used for efficient indexing into vectors and matrices. + https://petsc.org/release/manualpages/IS/IS/ + """ + @property + def dtype(self): + return CustomDtype('IS', modifier=' *') + + +class PointerIS(CallbackPointerIS): + @property + def _C_free(self): + destroy_calls = [ + petsc_call('ISDestroy', [Byref(self.indexify().subs({self.dim: i}))]) + for i in range(self._nindices) + ] + destroy_calls.append(petsc_call('PetscFree', [self.function])) + return destroy_calls + + +class CallbackPointerDM(PETScArrayObject): + @property + def dtype(self): + return CustomDtype('DM', modifier=' *') + + +class PointerDM(CallbackPointerDM): + @property + def _C_free(self): + destroy_calls = [ + petsc_call('DMDestroy', [Byref(self.indexify().subs({self.dim: i}))]) + for i in range(self._nindices) + ] + destroy_calls.append(petsc_call('PetscFree', [self.function])) + return destroy_calls + + +class PointerMat(PETScArrayObject): + _C_modifier = ' *' + + @property + def dtype(self): + return CustomDtype('Mat', modifier=' *') class ArgvSymbol(DataSymbol): diff --git a/devito/petsc/types/types.py b/devito/petsc/types/types.py index c3b43b84af..c7c3397181 100644 --- a/devito/petsc/types/types.py +++ b/devito/petsc/types/types.py @@ -29,20 +29,16 @@ class LinearSolveExpr(MetaData): needed to execute a linear solver. Linear problems are handled with `SNESSetType(snes, KSPONLY)`, enabling a unified interface for both linear and nonlinear solvers. - # TODO: extend this defaults: - 'ksp_type': String with the name of the PETSc Krylov method. Default is 'gmres' (Generalized Minimal Residual Method). https://petsc.org/main/manualpages/KSP/KSPType/ - - 'pc_type': String with the name of the PETSc preconditioner. Default is 'jacobi' (i.e diagonal scaling preconditioning). https://petsc.org/main/manualpages/PC/PCType/ - KSP tolerances: https://petsc.org/release/manualpages/KSP/KSPSetTolerances/ - - 'ksp_rtol': Relative convergence tolerance. Default is 1e-5. - 'ksp_atol': Absolute convergence for tolerance. Default @@ -55,8 +51,7 @@ class LinearSolveExpr(MetaData): """ __rargs__ = ('expr',) - __rkwargs__ = ('target', 'solver_parameters', 'matvecs', - 'formfuncs', 'formrhs', 'arrays', 'time_mapper', + __rkwargs__ = ('solver_parameters', 'fielddata', 'time_mapper', 'localinfo') defaults = { @@ -68,9 +63,8 @@ class LinearSolveExpr(MetaData): 'ksp_max_it': 1e4 # Maximum iterations } - def __new__(cls, expr, target=None, solver_parameters=None, - matvecs=None, formfuncs=None, formrhs=None, - arrays=None, time_mapper=None, localinfo=None, **kwargs): + def __new__(cls, expr, solver_parameters=None, + fielddata=None, time_mapper=None, localinfo=None, **kwargs): if solver_parameters is None: solver_parameters = cls.defaults @@ -82,12 +76,8 @@ def __new__(cls, expr, target=None, solver_parameters=None, obj = sympy.Function.__new__(cls, expr) obj._expr = expr - obj._target = target obj._solver_parameters = solver_parameters - obj._matvecs = matvecs - obj._formfuncs = formfuncs - obj._formrhs = formrhs - obj._arrays = arrays + obj._fielddata = fielddata if fielddata else FieldData() obj._time_mapper = time_mapper obj._localinfo = localinfo return obj @@ -105,21 +95,52 @@ def __hash__(self): def __eq__(self, other): return (isinstance(other, LinearSolveExpr) and - self.expr == other.expr and - self.target == other.target) + self.expr == other.expr) @property def expr(self): return self._expr @property - def target(self): - return self._target + def fielddata(self): + return self._fielddata @property def solver_parameters(self): return self._solver_parameters + @property + def time_mapper(self): + return self._time_mapper + + @property + def localinfo(self): + return self._localinfo + + @property + def grid(self): + return self.fielddata.grid + + @classmethod + def eval(cls, *args): + return None + + func = Reconstructable._rebuild + + +class FieldData: + def __init__(self, target=None, matvecs=None, formfuncs=None, formrhs=None, + arrays=None, **kwargs): + self._target = kwargs.get('target', target) + self._matvecs = matvecs + self._formfuncs = formfuncs + self._formrhs = formrhs + self._arrays = arrays + + @property + def target(self): + return self._target + @property def matvecs(self): return self._matvecs @@ -137,15 +158,160 @@ def arrays(self): return self._arrays @property - def time_mapper(self): - return self._time_mapper + def space_dimensions(self): + return self.target.space_dimensions @property - def localinfo(self): - return self._localinfo + def grid(self): + return self.target.grid - @classmethod - def eval(cls, *args): + @property + def space_order(self): + return self.target.space_order + + @property + def targets(self): + return (self.target,) + + +class MultipleFieldData(FieldData): + def __init__(self, submatrices=None): + self.field_data_list = [] + self._submatrices = submatrices + + def add_field_data(self, field_data): + self.field_data_list.append(field_data) + + def get_field_data(self, target): + for field_data in self.field_data_list: + if field_data.target == target: + return field_data + raise ValueError(f"FieldData with target {target} not found.") + pass + + @property + def target(self): return None - func = Reconstructable._rebuild + @property + def targets(self): + return tuple(field_data.target for field_data in self.field_data_list) + + @property + def space_dimensions(self): + space_dims = {field_data.space_dimensions for field_data in self.field_data_list} + if len(space_dims) > 1: + # TODO: This may not actually have to be the case, but enforcing it for now + raise ValueError( + "All targets within a PETScSolve have to have the same space dimensions." + ) + return space_dims.pop() + + @property + def grid(self): + grids = [t.grid for t in self.targets] + if len(set(grids)) > 1: + raise ValueError( + "All targets within a PETScSolve have to have the same grid." + ) + return grids.pop() + + @property + def space_order(self): + # NOTE: since we use DMDA to create vecs for the coupled solves, + # all fields must have the same space order + # ... re think this? limitation. For now, just force the + # space order to be the same. + # This isn't a problem for segregated solves. + space_orders = [t.space_order for t in self.targets] + if len(set(space_orders)) > 1: + raise ValueError( + "All targets within a PETScSolve have to have the same space order." + ) + return space_orders.pop() + + @property + def submatrices(self): + return self._submatrices + + +class SubMatrices: + def __init__(self, targets): + self.targets = targets + self.submatrices = self._initialize_submatrices() + + def _initialize_submatrices(self): + """ + Create a dict of submatrices for each target with metadata. + """ + submatrices = {} + num_targets = len(self.targets) + + for i, target in enumerate(self.targets): + submatrices[target] = {} + for j in range(num_targets): + key = f'J{i}{j}' + submatrices[target][key] = { + 'matvecs': None, + 'derivative_wrt': self.targets[j], + 'index': i * num_targets + j + } + + return submatrices + + @property + def submatrix_keys(self): + """ + Return a list of all submatrix keys (e.g., ['J00', 'J01', 'J10', 'J11']). + """ + return [key for submats in self.submatrices.values() for key in submats.keys()] + + @property + def nonzero_submatrix_keys(self): + """ + Returns a list of submats where 'matvecs' is not None. + """ + return [ + key + for submats in self.submatrices.values() + for key, value in submats.items() + if value['matvecs'] is not None + ] + + @property + def submat_to_index(self): + """ + Returns a dict mapping submatrix keys to their index. + """ + return { + key: value['index'] + for submats in self.submatrices.values() + for key, value in submats.items() + } + + def set_submatrix(self, field, key, matvecs): + """ + Set a specific submatrix for a field. + + Parameters + ---------- + field : Function + The target field that the submatrix operates on. + key: str + The identifier for the submatrix (e.g., 'J00', 'J01'). + matvecs: list of Eq + The matrix-vector equations forming the submatrix. + """ + if field in self.submatrices and key in self.submatrices[field]: + self.submatrices[field][key]["matvecs"] = matvecs + else: + raise KeyError(f'Invalid field ({field}) or submatrix key ({key})') + + def get_submatrix(self, field, key): + """ + Retrieve a specific submatrix. + """ + return self.submatrices.get(field, {}).get(key, None) + + def __repr__(self): + return str(self.submatrices) diff --git a/devito/symbolics/extended_sympy.py b/devito/symbolics/extended_sympy.py index dac1d20cb7..12000427e7 100644 --- a/devito/symbolics/extended_sympy.py +++ b/devito/symbolics/extended_sympy.py @@ -11,7 +11,7 @@ from devito.tools import (Pickable, Bunch, as_tuple, is_integer, float2, # noqa float3, float4, double2, double3, double4, int2, int3, int4, dtype_to_ctype, ctypes_to_cstr, ctypes_vector_mapper, - ctypes_to_cstr) + ctypes_to_cstr, CustomIntType) from devito.types import Symbol from devito.types.basic import Basic @@ -20,7 +20,8 @@ 'ListInitializer', 'Byref', 'IndexedPointer', 'Cast', 'DefFunction', 'MathFunction', 'InlineIf', 'ReservedWord', 'Keyword', 'String', 'Macro', 'Class', 'MacroArgument', 'Deref', 'Namespace', - 'Rvalue', 'Null', 'SizeOf', 'rfunc', 'BasicWrapperMixin', 'ValueLimit'] + 'Rvalue', 'Null', 'SizeOf', 'rfunc', 'BasicWrapperMixin', 'ValueLimit', + 'Mod'] class CondEq(sympy.Eq): @@ -90,9 +91,16 @@ def __new__(cls, lhs, rhs, params=None): # Perhaps it's a symbolic RHS -- but we wanna be sure it's of type int if not hasattr(rhs, 'dtype'): raise ValueError(f"Symbolic RHS `{rhs}` lacks dtype") - if not issubclass(rhs.dtype, np.integer): - raise ValueError(f"Symbolic RHS `{rhs}` must be of type `int`, found " - f"`{rhs.dtype}` instead") + + # TODO: Move into a utility function? + is_int_type = isinstance(rhs.dtype, type) and \ + issubclass(rhs.dtype, np.integer) + is_custom_int_type = isinstance(rhs.dtype, CustomIntType) + assert is_int_type or is_custom_int_type, ( + f"Symbolic RHS `{rhs}` must be of type `int`, " + f"found `{rhs.dtype}` instead" + ) + rhs = sympify(rhs) obj = sympy.Expr.__new__(cls, lhs, rhs) @@ -115,6 +123,26 @@ def __mul__(self, other): return super().__mul__(other) +class Mod(sympy.Expr): + # TODO: Add tests + is_Atom = True + is_commutative = True + + def __new__(cls, lhs, rhs, params=None): + rhs = sympify(rhs) + + obj = sympy.Expr.__new__(cls, lhs, rhs) + + obj.lhs = lhs + obj.rhs = rhs + return obj + + def __str__(self): + return "Mod(%s, %s)" % (self.lhs, self.rhs) + + __repr__ = __str__ + + class BasicWrapperMixin: """ @@ -167,7 +195,7 @@ def __new__(cls, call, pointer, params=None, **kwargs): pointer = Symbol(pointer) if isinstance(call, str): call = Symbol(call) - elif not isinstance(call, Basic): + elif not isinstance(call.base, Basic): raise ValueError("`call` must be a `devito.Basic` or a type " "with compatible interface") _params = [] diff --git a/devito/tools/dtypes_lowering.py b/devito/tools/dtypes_lowering.py index ce9813ce22..9394437905 100644 --- a/devito/tools/dtypes_lowering.py +++ b/devito/tools/dtypes_lowering.py @@ -15,7 +15,8 @@ 'double3', 'double4', 'dtypes_vector_mapper', 'dtype_to_mpidtype', 'dtype_to_cstr', 'dtype_to_ctype', 'infer_datasize', 'dtype_to_mpitype', 'dtype_len', 'ctypes_to_cstr', 'c_restrict_void_p', 'ctypes_vector_mapper', - 'is_external_ctype', 'infer_dtype', 'CustomDtype', 'mpi4py_mapper'] + 'is_external_ctype', 'infer_dtype', 'CustomDtype', 'mpi4py_mapper', + 'CustomIntType'] @@ -127,6 +128,11 @@ def __repr__(self): __str__ = __repr__ +# TODO: Consider if this should be an instance instead of a subclass? +class CustomIntType(CustomDtype): + pass + + # *** np.dtypes lowering @@ -313,6 +319,8 @@ def is_external_ctype(ctype, includes): True if `ctype` is known to be declared in one of the given `includes` files, False otherwise. """ + if isinstance(ctype, CustomDtype): + return False # Get the base type while issubclass(ctype, ctypes._Pointer): ctype = ctype._type_ diff --git a/devito/types/array.py b/devito/types/array.py index 105cdd21da..42413f9534 100644 --- a/devito/types/array.py +++ b/devito/types/array.py @@ -60,6 +60,21 @@ def shape_allocated(self): def is_const(self): return self._is_const + @property + def _C_free(self): + """ + A symbolic destructor for the Array, injected in the generated code. + + Notes + ----- + To be overridden by subclasses, ignored otherwise. + """ + return None + + @property + def _C_free_priority(self): + return 0 + class Array(ArrayBasic): diff --git a/devito/types/object.py b/devito/types/object.py index 1668f2c289..6eea49f6a1 100644 --- a/devito/types/object.py +++ b/devito/types/object.py @@ -1,14 +1,14 @@ from ctypes import byref import sympy -from devito.tools import Pickable, as_tuple, sympy_mutex +from devito.tools import Pickable, as_tuple, sympy_mutex, CustomDtype from devito.types.args import ArgProvider from devito.types.caching import Uncached from devito.types.basic import Basic, LocalType from devito.types.utils import CtypesFactory -__all__ = ['Object', 'LocalObject', 'CompositeObject', 'CCompositeObject'] +__all__ = ['Object', 'LocalObject', 'CompositeObject', 'LocalCompositeObject'] class AbstractObject(Basic, sympy.Basic, Pickable): @@ -242,19 +242,30 @@ def _mem_global(self): return self._is_global -class CCompositeObject(CompositeObject, LocalType): +class LocalCompositeObject(CompositeObject, LocalType): """ Object with composite type (e.g., a C struct) defined in C. """ - __rargs__ = ('name', 'pname', 'pfields') + __rargs__ = ('name', 'pname', 'fields') - def __init__(self, name, pname, pfields, liveness='lazy'): - super().__init__(name, pname, pfields) + def __init__(self, name, pname, fields, modifier=None, liveness='lazy'): + dtype = CustomDtype(f"struct {pname}", modifier=modifier) + Object.__init__(self, name, dtype, None) + self._pname = pname assert liveness in ['eager', 'lazy'] self._liveness = liveness + self._fields = fields @property - def dtype(self): - return self._dtype._type_ + def fields(self): + return self._fields + + @property + def _fields_(self): + return [(i._C_name, i._C_ctype) for i in self.fields] + + @property + def __name__(self): + return self.pname diff --git a/tests/test_petsc.py b/tests/test_petsc.py index dafc22211c..1f0f8bc91e 100644 --- a/tests/test_petsc.py +++ b/tests/test_petsc.py @@ -4,14 +4,14 @@ from conftest import skipif from devito import (Grid, Function, TimeFunction, Eq, Operator, switchconfig, - configuration) + configuration, norm) from devito.ir.iet import (Call, ElementalFunction, Definition, DummyExpr, FindNodes, retrieve_iteration_tree) -from devito.types import Constant, CCompositeObject +from devito.types import Constant, LocalCompositeObject from devito.passes.iet.languages.C import CDataManager from devito.petsc.types import (DM, Mat, LocalVec, PetscMPIInt, KSP, PC, KSPConvergedReason, PETScArray, - LinearSolveExpr) + LinearSolveExpr, FieldData, MultipleFieldData) from devito.petsc.solve import PETScSolve, separate_eqn, centre_stencil from devito.petsc.iet.nodes import Expression from devito.petsc.initialize import PetscInitialize @@ -126,20 +126,21 @@ def test_petsc_solve(): callable_roots = [meta_call.root for meta_call in op._func_table.values()] - matvec_callback = [root for root in callable_roots if root.name == 'MyMatShellMult_0'] + matvec_callback = [root for root in callable_roots if root.name == 'MatMult0'] - formrhs_callback = [root for root in callable_roots if root.name == 'FormRHS_0'] + formrhs_callback = [root for root in callable_roots if root.name == 'FormRHS0'] action_expr = FindNodes(Expression).visit(matvec_callback[0]) rhs_expr = FindNodes(Expression).visit(formrhs_callback[0]) - assert str(action_expr[-1].expr.rhs) == \ - 'x_matvec_f[x + 1, y + 2]/lctx->h_x**2' + \ - ' - 2.0*x_matvec_f[x + 2, y + 2]/lctx->h_x**2' + \ - ' + x_matvec_f[x + 3, y + 2]/lctx->h_x**2' + \ - ' + x_matvec_f[x + 2, y + 1]/lctx->h_y**2' + \ - ' - 2.0*x_matvec_f[x + 2, y + 2]/lctx->h_y**2' + \ - ' + x_matvec_f[x + 2, y + 3]/lctx->h_y**2' + assert str(action_expr[-1].expr.rhs) == ( + 'x_f[x + 1, y + 2]/ctx0->h_x**2' + ' - 2.0*x_f[x + 2, y + 2]/ctx0->h_x**2' + ' + x_f[x + 3, y + 2]/ctx0->h_x**2' + ' + x_f[x + 2, y + 1]/ctx0->h_y**2' + ' - 2.0*x_f[x + 2, y + 2]/ctx0->h_y**2' + ' + x_f[x + 2, y + 3]/ctx0->h_y**2' + ) assert str(rhs_expr[-1].expr.rhs) == 'g[x + 2, y + 2]' @@ -213,12 +214,12 @@ def test_petsc_cast(): cb2 = [meta_call.root for meta_call in op2._func_table.values()] cb3 = [meta_call.root for meta_call in op3._func_table.values()] - assert 'float (*restrict x_matvec_f1) = ' + \ - '(float (*)) x_matvec_f1_vec;' in str(cb1[0]) - assert 'float (*restrict x_matvec_f2)[info.gxm] = ' + \ - '(float (*)[info.gxm]) x_matvec_f2_vec;' in str(cb2[0]) - assert 'float (*restrict x_matvec_f3)[info.gym][info.gxm] = ' + \ - '(float (*)[info.gym][info.gxm]) x_matvec_f3_vec;' in str(cb3[0]) + assert 'float (*restrict x_f1) = ' + \ + '(float (*)) x_f1_vec;' in str(cb1[0]) + assert 'float (*restrict x_f2)[info.gxm] = ' + \ + '(float (*)[info.gxm]) x_f2_vec;' in str(cb2[0]) + assert 'float (*restrict x_f3)[info.gym][info.gxm] = ' + \ + '(float (*)[info.gym][info.gxm]) x_f3_vec;' in str(cb3[0]) @skipif('petsc') @@ -233,8 +234,9 @@ def test_LinearSolveExpr(): linsolveexpr = LinearSolveExpr(eqn.rhs, target=f) - # Check the target - assert linsolveexpr.target == f + # TODO: maybe expand this test now to check the fielddata etc + linsolveexpr = LinearSolveExpr(eqn.rhs) + # Check the solver parameters assert linsolveexpr.solver_parameters == \ {'ksp_type': 'gmres', 'pc_type': 'jacobi', 'ksp_rtol': 1e-05, @@ -266,15 +268,15 @@ def test_dmda_create(): op3 = Operator(petsc3, opt='noop') assert 'PetscCall(DMDACreate1d(PETSC_COMM_WORLD,DM_BOUNDARY_GHOSTED,' + \ - '2,1,2,NULL,&(da_0)));' in str(op1) + '2,1,2,NULL,&(da0)));' in str(op1) assert 'PetscCall(DMDACreate2d(PETSC_COMM_WORLD,DM_BOUNDARY_GHOSTED,' + \ - 'DM_BOUNDARY_GHOSTED,DMDA_STENCIL_BOX,2,2,1,1,1,4,NULL,NULL,&(da_0)));' \ + 'DM_BOUNDARY_GHOSTED,DMDA_STENCIL_BOX,2,2,1,1,1,4,NULL,NULL,&(da0)));' \ in str(op2) assert 'PetscCall(DMDACreate3d(PETSC_COMM_WORLD,DM_BOUNDARY_GHOSTED,' + \ 'DM_BOUNDARY_GHOSTED,DM_BOUNDARY_GHOSTED,DMDA_STENCIL_BOX,6,5,4' + \ - ',1,1,1,1,6,NULL,NULL,NULL,&(da_0)));' in str(op3) + ',1,1,1,1,6,NULL,NULL,NULL,&(da0)));' in str(op3) @skipif('petsc') @@ -301,9 +303,9 @@ def test_cinterface_petsc_struct(): assert 'include "%s.h"' % name in ccode - # The public `struct MatContext` only appears in the header file - assert 'struct J_0_ctx\n{' not in ccode - assert 'struct J_0_ctx\n{' in hcode + # The public `struct UserCtx` only appears in the header file + assert 'struct UserCtx0\n{' not in ccode + assert 'struct UserCtx0\n{' in hcode @skipif('petsc') @@ -571,14 +573,14 @@ def test_callback_arguments(): with switchconfig(openmp=False): op = Operator(petsc1) - mv = op._func_table['MyMatShellMult_0'].root - ff = op._func_table['FormFunction_0'].root + mv = op._func_table['MatMult0'].root + ff = op._func_table['FormFunction0'].root assert len(mv.parameters) == 3 assert len(ff.parameters) == 4 - assert str(mv.parameters) == '(J_0, X_global_0, Y_global_0)' - assert str(ff.parameters) == '(snes_0, X_global_0, F_global_0, dummy)' + assert str(mv.parameters) == '(J, X, Y)' + assert str(ff.parameters) == '(snes, X, F, dummy)' @skipif('petsc') @@ -611,7 +613,7 @@ def test_petsc_struct(): assert mu2 in op.parameters # Check PETSc struct not in op.parameters - assert all(not isinstance(i, CCompositeObject) for i in op.parameters) + assert all(not isinstance(i, LocalCompositeObject) for i in op.parameters) @skipif('petsc') @@ -630,8 +632,7 @@ def test_apply(): # Build the op op = Operator(petsc) - # Check the Operator runs without errors. Not verifying output for - # now. Need to consolidate BC implementation + # Check the Operator runs without errors op.apply() # Verify that users can override `mu` @@ -656,11 +657,11 @@ def test_petsc_frees(): frees = op.body.frees # Check the frees appear in the following order - assert str(frees[0]) == 'PetscCall(VecDestroy(&(b_global_0)));' - assert str(frees[1]) == 'PetscCall(VecDestroy(&(x_global_0)));' - assert str(frees[2]) == 'PetscCall(MatDestroy(&(J_0)));' - assert str(frees[3]) == 'PetscCall(SNESDestroy(&(snes_0)));' - assert str(frees[4]) == 'PetscCall(DMDestroy(&(da_0)));' + assert str(frees[0]) == 'PetscCall(VecDestroy(&(bglobal0)));' + assert str(frees[1]) == 'PetscCall(VecDestroy(&(xglobal0)));' + assert str(frees[2]) == 'PetscCall(MatDestroy(&(J0)));' + assert str(frees[3]) == 'PetscCall(SNESDestroy(&(snes0)));' + assert str(frees[4]) == 'PetscCall(DMDestroy(&(da0)));' @skipif('petsc') @@ -679,8 +680,8 @@ def test_calls_to_callbacks(): ccode = str(op.ccode) - assert '(void (*)(void))MyMatShellMult_0' in ccode - assert 'PetscCall(SNESSetFunction(snes_0,NULL,FormFunction_0,NULL));' in ccode + assert '(void (*)(void))MatMult0' in ccode + assert 'PetscCall(SNESSetFunction(snes0,NULL,FormFunction0,(void*)(da0)));' in ccode @skipif('petsc') @@ -702,7 +703,7 @@ def test_start_ptr(): op1 = Operator(petsc1) # Verify the case with modulo time stepping - assert 'float * start_ptr_0 = t1*localsize_0 + (float*)(u1_vec->data);' in str(op1) + assert 'float * u1_ptr0 = t1*localsize0 + (float*)(u1_vec->data);' in str(op1) # Verify the case with no modulo time stepping u2 = TimeFunction(name='u2', grid=grid, space_order=2, dtype=np.float32, save=5) @@ -712,7 +713,7 @@ def test_start_ptr(): with switchconfig(openmp=False): op2 = Operator(petsc2) - assert 'float * start_ptr_0 = (time + 1)*localsize_0 + ' + \ + assert 'float * u2_ptr0 = (time + 1)*localsize0 + ' + \ '(float*)(u2_vec->data);' in str(op2) @@ -735,12 +736,12 @@ def test_time_loop(): with switchconfig(openmp=False): op1 = Operator(petsc1) body1 = str(op1.body) - rhs1 = str(op1._func_table['FormRHS_0'].root.ccode) + rhs1 = str(op1._func_table['FormRHS0'].root.ccode) assert 'ctx0.t0 = t0' in body1 assert 'ctx0.t1 = t1' not in body1 - assert 'lctx->t0' in rhs1 - assert 'lctx->t1' not in rhs1 + assert 'ctx0->t0' in rhs1 + assert 'ctx0->t1' not in rhs1 # Non-modulo time stepping u2 = TimeFunction(name='u2', grid=grid, space_order=2, save=5) @@ -750,10 +751,10 @@ def test_time_loop(): with switchconfig(openmp=False): op2 = Operator(petsc2) body2 = str(op2.body) - rhs2 = str(op2._func_table['FormRHS_0'].root.ccode) + rhs2 = str(op2._func_table['FormRHS0'].root.ccode) assert 'ctx0.time = time' in body2 - assert 'lctx->time' in rhs2 + assert 'ctx0->time' in rhs2 # Modulo time stepping with more than one time step # used in one of the callback functions @@ -762,12 +763,12 @@ def test_time_loop(): with switchconfig(openmp=False): op3 = Operator(petsc3) body3 = str(op3.body) - rhs3 = str(op3._func_table['FormRHS_0'].root.ccode) + rhs3 = str(op3._func_table['FormRHS0'].root.ccode) assert 'ctx0.t0 = t0' in body3 assert 'ctx0.t1 = t1' in body3 - assert 'lctx->t0' in rhs3 - assert 'lctx->t1' in rhs3 + assert 'ctx0->t0' in rhs3 + assert 'ctx0->t1' in rhs3 # Multiple petsc solves within the same time loop v2 = Function(name='v2', grid=grid, space_order=2) @@ -781,3 +782,245 @@ def test_time_loop(): assert 'ctx0.t0 = t0' in body4 assert body4.count('ctx0.t0 = t0') == 1 + + +class TestCoupledLinear: + # The coupled interface can be used even for uncoupled problems, meaning + # the equations will be solved within a single matrix system. + # These tests use simple problems to validate functionality, but they help + # ensure correctness in code generation. + # TODO: Add more comprehensive tests for fully coupled problems. + # TODO: Add subdomain tests, time loop, multiple coupled etc. + + @skipif('petsc') + def test_coupled_vs_non_coupled(self): + grid = Grid(shape=(11, 11), dtype=np.float64) + + functions = [Function(name=n, grid=grid, space_order=2, dtype=np.float64) + for n in ['e', 'f', 'g', 'h']] + e, f, g, h = functions + + f.data[:] = 5. + h.data[:] = 5. + + eq1 = Eq(e.laplace, f) + eq2 = Eq(g.laplace, h) + + # Non-coupled + petsc1 = PETScSolve(eq1, target=e) + petsc2 = PETScSolve(eq2, target=g) + + with switchconfig(openmp=False): + op1 = Operator(petsc1 + petsc2, opt='noop') + op1.apply() + + enorm1 = norm(e) + gnorm1 = norm(g) + + # Reset + e.data[:] = 0 + g.data[:] = 0 + + # Coupled + # TODO: Need more friendly API for coupled - just + # using a dict for now + petsc3 = PETScSolve({e: [eq1], g: [eq2]}) + with switchconfig(openmp=False): + op2 = Operator(petsc3, opt='noop') + op2.apply() + + enorm2 = norm(e) + gnorm2 = norm(g) + + print('enorm1:', enorm1) + print('enorm2:', enorm2) + assert np.isclose(enorm1, enorm2, rtol=1e-16) + assert np.isclose(gnorm1, gnorm2, rtol=1e-16) + + callbacks1 = [meta_call.root for meta_call in op1._func_table.values()] + callbacks2 = [meta_call.root for meta_call in op2._func_table.values()] + + # Solving for multiple fields within the same matrix system requires + # additional machinery and more callback functions + assert len(callbacks1) == 8 + assert len(callbacks2) == 11 + + # Check fielddata type + fielddata1 = petsc1[0].rhs.fielddata + fielddata2 = petsc2[0].rhs.fielddata + fielddata3 = petsc3[0].rhs.fielddata + + assert isinstance(fielddata1, FieldData) + assert isinstance(fielddata2, FieldData) + assert isinstance(fielddata3, MultipleFieldData) + + @skipif('petsc') + def test_coupled_structs(self): + grid = Grid(shape=(11, 11)) + + functions = [Function(name=n, grid=grid, space_order=2, dtype=np.float64) + for n in ['e', 'f', 'g', 'h']] + e, f, g, h = functions + + eq1 = Eq(e + 5, f) + eq2 = Eq(g + 10, h) + + petsc = PETScSolve({f: [eq1], h: [eq2]}) + + name = "foo" + with switchconfig(openmp=False): + op = Operator(petsc, name=name) + + # Trigger the generation of a .c and a .h files + ccode, hcode = op.cinterface(force=True) + + dirname = op._compiler.get_jit_dir() + assert os.path.isfile(os.path.join(dirname, f"{name}.c")) + assert os.path.isfile(os.path.join(dirname, f"{name}.h")) + + ccode = str(ccode) + hcode = str(hcode) + + assert f'include "{name}.h"' in ccode + + # The public `struct JacobianCtx` only appears in the header file + assert 'struct JacobianCtx\n{' not in ccode + assert 'struct JacobianCtx\n{' in hcode + + # The public `struct SubMatrixCtx` only appears in the header file + assert 'struct SubMatrixCtx\n{' not in ccode + assert 'struct SubMatrixCtx\n{' in hcode + + # The public `struct UserCtx0` only appears in the header file + assert 'struct UserCtx0\n{' not in ccode + assert 'struct UserCtx0\n{' in hcode + + @skipif('petsc') + def test_coupled_frees(self): + grid = Grid(shape=(11, 11), dtype=np.float64) + + functions = [Function(name=n, grid=grid, space_order=2, dtype=np.float64) + for n in ['e', 'f', 'g', 'h']] + e, f, g, h = functions + + eq1 = Eq(e.laplace, h) + eq2 = Eq(f.laplace, h) + eq3 = Eq(g.laplace, h) + + petsc1 = PETScSolve({e: [eq1], f: [eq2]}) + petsc2 = PETScSolve({e: [eq1], f: [eq2], g: [eq3]}) + + with switchconfig(openmp=False): + op1 = Operator(petsc1, opt='noop') + op2 = Operator(petsc2, opt='noop') + + frees1 = op1.body.frees + frees2 = op2.body.frees + + # Check solver with two fields + # IS destroys + assert str(frees1[0]) == 'PetscCall(ISDestroy(&(fields0[0])));' + assert str(frees1[1]) == 'PetscCall(ISDestroy(&(fields0[1])));' + assert str(frees1[2]) == 'PetscCall(PetscFree(fields0));' + # Sub DM destroys + assert str(frees1[3]) == 'PetscCall(DMDestroy(&(subdms0[0])));' + assert str(frees1[4]) == 'PetscCall(DMDestroy(&(subdms0[1])));' + assert str(frees1[5]) == 'PetscCall(PetscFree(subdms0));' + + # Check solver with three fields + # IS destroys + assert str(frees2[0]) == 'PetscCall(ISDestroy(&(fields0[0])));' + assert str(frees2[1]) == 'PetscCall(ISDestroy(&(fields0[1])));' + assert str(frees2[2]) == 'PetscCall(ISDestroy(&(fields0[2])));' + assert str(frees2[3]) == 'PetscCall(PetscFree(fields0));' + # Sub DM destroys + assert str(frees2[4]) == 'PetscCall(DMDestroy(&(subdms0[0])));' + assert str(frees2[5]) == 'PetscCall(DMDestroy(&(subdms0[1])));' + assert str(frees2[6]) == 'PetscCall(DMDestroy(&(subdms0[2])));' + assert str(frees2[7]) == 'PetscCall(PetscFree(subdms0));' + + @skipif('petsc') + def test_dmda_dofs(self): + grid = Grid(shape=(11, 11)) + + functions = [Function(name=n, grid=grid, space_order=2, dtype=np.float64) + for n in ['e', 'f', 'g', 'h']] + e, f, g, h = functions + + eq1 = Eq(e.laplace, h) + eq2 = Eq(f.laplace, h) + eq3 = Eq(g.laplace, h) + + petsc1 = PETScSolve({e: [eq1]}) + petsc2 = PETScSolve({e: [eq1], f: [eq2]}) + petsc3 = PETScSolve({e: [eq1], f: [eq2], g: [eq3]}) + + with switchconfig(openmp=False): + op1 = Operator(petsc1, opt='noop') + op2 = Operator(petsc2, opt='noop') + op3 = Operator(petsc3, opt='noop') + + # Check the number of dofs in the DMDA for each field + assert 'PetscCall(DMDACreate2d(PETSC_COMM_WORLD,DM_BOUNDARY_GHOSTED,' + \ + 'DM_BOUNDARY_GHOSTED,DMDA_STENCIL_BOX,11,11,1,1,1,2,NULL,NULL,&(da0)));' \ + in str(op1) + + assert 'PetscCall(DMDACreate2d(PETSC_COMM_WORLD,DM_BOUNDARY_GHOSTED,' + \ + 'DM_BOUNDARY_GHOSTED,DMDA_STENCIL_BOX,11,11,1,1,2,2,NULL,NULL,&(da0)));' \ + in str(op2) + + assert 'PetscCall(DMDACreate2d(PETSC_COMM_WORLD,DM_BOUNDARY_GHOSTED,' + \ + 'DM_BOUNDARY_GHOSTED,DMDA_STENCIL_BOX,11,11,1,1,3,2,NULL,NULL,&(da0)));' \ + in str(op3) + + @skipif('petsc') + def test_submatrices(self): + grid = Grid(shape=(11, 11)) + + functions = [Function(name=n, grid=grid, space_order=2, dtype=np.float64) + for n in ['e', 'f', 'g', 'h']] + e, f, g, h = functions + + eq1 = Eq(e.laplace, f) + eq2 = Eq(g.laplace, h) + + petsc = PETScSolve({e: [eq1], g: [eq2]}) + + submatrices = petsc[0].rhs.fielddata.submatrices + + j00 = submatrices.get_submatrix(e, 'J00') + j01 = submatrices.get_submatrix(e, 'J01') + j10 = submatrices.get_submatrix(g, 'J10') + j11 = submatrices.get_submatrix(g, 'J11') + + # Check the number of submatrices + assert len(submatrices.submatrix_keys) == 4 + assert str(submatrices.submatrix_keys) == "['J00', 'J01', 'J10', 'J11']" + + # Technically a non-coupled problem, so the only non-zero submatrices + # should be the diagonal ones i.e J00 and J11 + assert submatrices.nonzero_submatrix_keys == ['J00', 'J11'] + assert submatrices.get_submatrix(e, 'J01')['matvecs'] is None + assert submatrices.get_submatrix(g, 'J10')['matvecs'] is None + + j00 = submatrices.get_submatrix(e, 'J00') + j11 = submatrices.get_submatrix(g, 'J11') + + assert str(j00['matvecs'][0]) == 'Eq(y_e(x, y),' \ + + ' Derivative(x_e(x, y), (x, 2)) + Derivative(x_e(x, y), (y, 2)))' + + assert str(j11['matvecs'][0]) == 'Eq(y_g(x, y),' \ + + ' Derivative(x_g(x, y), (x, 2)) + Derivative(x_g(x, y), (y, 2)))' + + # Check the derivative wrt fields + assert j00['derivative_wrt'] == e + assert j01['derivative_wrt'] == g + assert j10['derivative_wrt'] == e + assert j11['derivative_wrt'] == g + + # TODO: + # @skipif('petsc') + # def test_create_submats(self): + + # add tests for all new callbacks + # def test_create_whole_matmult(): From 1146a07f29ed9b8db590300fbf4800dbe00f5978 Mon Sep 17 00:00:00 2001 From: ZoeLeibowitz Date: Fri, 21 Mar 2025 16:51:28 +0000 Subject: [PATCH 05/84] misc: Rebase fixes --- .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/petsc/iet/routines.py | 13 ++++++++----- devito/petsc/types/object.py | 8 +++----- devito/tools/dtypes_lowering.py | 1 - 10 files changed, 25 insertions(+), 11 deletions(-) diff --git a/.github/workflows/examples-mpi.yml b/.github/workflows/examples-mpi.yml index e68f250908..e44a03410d 100644 --- a/.github/workflows/examples-mpi.yml +++ b/.github/workflows/examples-mpi.yml @@ -17,9 +17,11 @@ on: push: branches: - main + - master pull_request: branches: - main + - master jobs: build: diff --git a/.github/workflows/examples.yml b/.github/workflows/examples.yml index 902759f599..9b6d5a9ade 100644 --- a/.github/workflows/examples.yml +++ b/.github/workflows/examples.yml @@ -10,9 +10,11 @@ on: push: branches: - main + - master pull_request: branches: - main + - master jobs: tutorials: diff --git a/.github/workflows/flake8.yml b/.github/workflows/flake8.yml index 3a645541f7..40b57d7706 100644 --- a/.github/workflows/flake8.yml +++ b/.github/workflows/flake8.yml @@ -10,9 +10,11 @@ on: push: branches: - main + - master pull_request: branches: - main + - master jobs: flake8: diff --git a/.github/workflows/pytest-core-mpi.yml b/.github/workflows/pytest-core-mpi.yml index 17b329a4a8..6de5933032 100644 --- a/.github/workflows/pytest-core-mpi.yml +++ b/.github/workflows/pytest-core-mpi.yml @@ -10,9 +10,11 @@ on: push: branches: - main + - master pull_request: branches: - main + - master jobs: test-mpi-basic: diff --git a/.github/workflows/pytest-core-nompi.yml b/.github/workflows/pytest-core-nompi.yml index 666b365467..d7c508cec2 100644 --- a/.github/workflows/pytest-core-nompi.yml +++ b/.github/workflows/pytest-core-nompi.yml @@ -10,9 +10,11 @@ on: push: branches: - main + - master pull_request: branches: - main + - master jobs: pytest: diff --git a/.github/workflows/pytest-petsc.yml b/.github/workflows/pytest-petsc.yml index 91483e2fc4..090748b222 100644 --- a/.github/workflows/pytest-petsc.yml +++ b/.github/workflows/pytest-petsc.yml @@ -9,9 +9,11 @@ on: # but only for the master branch push: branches: + - main - master pull_request: branches: + - main - master jobs: diff --git a/.github/workflows/tutorials.yml b/.github/workflows/tutorials.yml index a4b1c2fc85..c6e9c2adba 100644 --- a/.github/workflows/tutorials.yml +++ b/.github/workflows/tutorials.yml @@ -10,9 +10,11 @@ on: push: branches: - main + - master pull_request: branches: - main + - master jobs: tutorials: diff --git a/devito/petsc/iet/routines.py b/devito/petsc/iet/routines.py index a70e987e71..4c28231cec 100644 --- a/devito/petsc/iet/routines.py +++ b/devito/petsc/iet/routines.py @@ -4,7 +4,7 @@ from devito.ir.iet import (Call, FindSymbols, List, Uxreplace, CallableBody, Dereference, DummyExpr, BlankLine, Callable, FindNodes, retrieve_iteration_tree, filter_iterations, Iteration) -from devito.symbolics import (Byref, FieldFromPointer, cast_mapper, VOIDP, +from devito.symbolics import (Byref, FieldFromPointer, cast, VOID, FieldFromComposite, IntDiv, Deref, Mod) from devito.symbolics.unevaluation import Mul from devito.types.basic import AbstractFunction @@ -671,8 +671,10 @@ def _whole_formfunc_body(self): calls += ( petsc_call('VecGetSubVector', [X, field_ptr, Byref(sobjs[x_name])]), petsc_call('VecGetSubVector', [F, field_ptr, Byref(sobjs[f_name])]), - petsc_call(self.formfuncs[i].name, [objs['snes'], sobjs[x_name], - sobjs[f_name], VOIDP(objs['LocalSubdms'].indexed[i])]), + petsc_call(self.formfuncs[i].name, [ + objs['snes'], sobjs[x_name], sobjs[f_name], + VOID(objs['LocalSubdms'].indexed[i], stars='*') + ]), petsc_call('VecRestoreSubVector', [X, field_ptr, Byref(sobjs[x_name])]), petsc_call('VecRestoreSubVector', [F, field_ptr, Byref(sobjs[f_name])]), ) @@ -992,7 +994,7 @@ def snes_ctx(self): The [optional] context for private data for the function evaluation routine. https://petsc.org/main/manualpages/SNES/SNESSetFunction/ """ - return VOIDP(self.solver_objs['dmda']) + return VOID(self.solver_objs['dmda'], stars='*') def _setup(self): objs = self.objs @@ -1526,11 +1528,12 @@ def replace_array(self, target): xlocal = sobjs.get(f'xlocal{target.name}', sobjs['xlocal']) start_ptr = sobjs[f'{target.name}_ptr'] + caster = cast(target.dtype, '*') return ( petsc_call('VecGetSize', [xlocal, Byref(sobjs['localsize'])]), DummyExpr( start_ptr, - cast_mapper[(target.dtype, '*')]( + caster( FieldFromPointer(target._C_field_data, target._C_symbol) ) + Mul(target_time, sobjs['localsize']), init=True diff --git a/devito/petsc/types/object.py b/devito/petsc/types/object.py index 9acf7def46..bf4339377d 100644 --- a/devito/petsc/types/object.py +++ b/devito/petsc/types/object.py @@ -2,7 +2,7 @@ from devito.tools import CustomDtype, dtype_to_cstr, as_tuple, CustomIntType from devito.types import (LocalObject, LocalCompositeObject, ModuloDimension, TimeDimension, ArrayObject, CustomDimension) -from devito.symbolics import Byref, Cast +from devito.symbolics import Byref, cast from devito.types.basic import DataSymbol from devito.petsc.iet.utils import petsc_call @@ -42,8 +42,7 @@ def _C_free_priority(self): return 4 -class DMCast(Cast): - _base_typ = 'DM' +DMCast = cast('DM') class CallbackMat(LocalObject): @@ -232,8 +231,7 @@ def __init__(self, name='subctx', pname='SubMatrixCtx', fields=None, _C_modifier = None -class JacobianStructCast(Cast): - _base_typ = 'struct JacobianCtx *' +JacobianStructCast = cast('struct JacobianCtx *') class PETScArrayObject(ArrayObject): diff --git a/devito/tools/dtypes_lowering.py b/devito/tools/dtypes_lowering.py index 9394437905..c78f85e1e2 100644 --- a/devito/tools/dtypes_lowering.py +++ b/devito/tools/dtypes_lowering.py @@ -19,7 +19,6 @@ 'CustomIntType'] - # *** Custom np.dtypes # NOTE: the following is inspired by pyopencl.cltypes From cccb988a042de47803c700a06d9c45450a34aa3e Mon Sep 17 00:00:00 2001 From: ZoeLeibowitz <90093761+ZoeLeibowitz@users.noreply.github.com> Date: Mon, 24 Mar 2025 13:45:00 +0000 Subject: [PATCH 06/84] Fix sympy_dtype for expressions containing FieldFromPointer objects (#66) --- devito/symbolics/inspection.py | 14 +++++++++----- 1 file changed, 9 insertions(+), 5 deletions(-) diff --git a/devito/symbolics/inspection.py b/devito/symbolics/inspection.py index 709487ffbe..522cd8a379 100644 --- a/devito/symbolics/inspection.py +++ b/devito/symbolics/inspection.py @@ -307,12 +307,16 @@ def sympy_dtype(expr, base=None, default=None, smin=None): if expr is None: return default - # TODO: Edit/fix/update according to PR #2513 + def inspect_args(e, dtypes): + for arg in e.args: + dtype = getattr(arg, "dtype", None) + if dtype is not None: + dtypes.add(dtype) + else: + inspect_args(arg, dtypes) + dtypes = {base} - {None} - for i in expr.args: - dtype = getattr(i, 'dtype', None) - if dtype: - dtypes.add(dtype) + inspect_args(expr, dtypes) dtype = infer_dtype(dtypes) From f61d4dd27e046a7866c82d5d0b0315c5088f324b Mon Sep 17 00:00:00 2001 From: ZoeLeibowitz <90093761+ZoeLeibowitz@users.noreply.github.com> Date: Wed, 26 Mar 2025 22:33:17 +0000 Subject: [PATCH 07/84] compiler: Add PETSc language and printer (#67) --- conftest.py | 17 +- devito/core/__init__.py | 5 +- devito/core/cpu.py | 23 ++- devito/operator/registry.py | 2 +- devito/passes/iet/languages/C.py | 7 + devito/passes/iet/languages/targets.py | 9 +- devito/petsc/iet/passes.py | 8 +- devito/petsc/initialize.py | 19 +-- devito/petsc/types/array.py | 4 - devito/petsc/types/object.py | 4 +- devito/petsc/types/types.py | 10 ++ devito/petsc/utils.py | 88 +++++++---- examples/petsc/petsc_test.py | 10 +- tests/test_petsc.py | 211 +++++++++++++------------ 14 files changed, 248 insertions(+), 169 deletions(-) diff --git a/conftest.py b/conftest.py index fc6090c604..c0339f722b 100644 --- a/conftest.py +++ b/conftest.py @@ -14,7 +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 +from devito.petsc.utils import PetscOSError, get_petsc_dir try: from mpi4py import MPI # noqa @@ -95,18 +95,11 @@ def skipif(items, whole_module=False): 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" + try: + _ = get_petsc_dir() + except PetscOSError: + skipit = "PETSc is not installed" 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/core/__init__.py b/devito/core/__init__.py index 5b7ad7878b..32d391e818 100644 --- a/devito/core/__init__.py +++ b/devito/core/__init__.py @@ -6,7 +6,8 @@ Cpu64CustomOperator, Cpu64CustomCXXOperator, Cpu64CXXNoopCOperator, Cpu64CXXNoopOmpOperator, Cpu64AdvCXXOperator, Cpu64AdvCXXOmpOperator, - Cpu64FsgCXXOperator, Cpu64FsgCXXOmpOperator + Cpu64FsgCXXOperator, Cpu64FsgCXXOmpOperator, + Cpu64NoopPetscOperator, Cpu64AdvPetscOperator ) from devito.core.intel import ( Intel64AdvCOperator, Intel64AdvOmpOperator, @@ -40,11 +41,13 @@ operator_registry.add(Cpu64NoopOmpOperator, Cpu64, 'noop', 'openmp') operator_registry.add(Cpu64CXXNoopCOperator, Cpu64, 'noop', 'CXX') operator_registry.add(Cpu64CXXNoopOmpOperator, Cpu64, 'noop', 'CXXopenmp') +operator_registry.add(Cpu64NoopPetscOperator, Cpu64, 'noop', 'petsc') operator_registry.add(Cpu64AdvCOperator, Cpu64, 'advanced', 'C') operator_registry.add(Cpu64AdvOmpOperator, Cpu64, 'advanced', 'openmp') operator_registry.add(Cpu64AdvCXXOperator, Cpu64, 'advanced', 'CXX') operator_registry.add(Cpu64AdvCXXOmpOperator, Cpu64, 'advanced', 'CXXopenmp') +operator_registry.add(Cpu64AdvPetscOperator, Cpu64, 'advanced', 'petsc') operator_registry.add(Cpu64FsgCOperator, Cpu64, 'advanced-fsg', 'C') operator_registry.add(Cpu64FsgOmpOperator, Cpu64, 'advanced-fsg', 'openmp') diff --git a/devito/core/cpu.py b/devito/core/cpu.py index d96d03e31d..cf225305f9 100644 --- a/devito/core/cpu.py +++ b/devito/core/cpu.py @@ -11,13 +11,14 @@ from devito.passes.iet import (CTarget, CXXTarget, COmpTarget, CXXOmpTarget, avoid_denormals, linearize, mpiize, hoist_prodders, relax_incr_dimensions, - check_stability) + check_stability, PetscTarget) from devito.tools import timed_pass __all__ = ['Cpu64NoopCOperator', 'Cpu64NoopOmpOperator', 'Cpu64AdvCOperator', 'Cpu64AdvOmpOperator', 'Cpu64FsgCOperator', 'Cpu64FsgOmpOperator', 'Cpu64CustomOperator', 'Cpu64CustomCXXOperator', 'Cpu64AdvCXXOperator', - 'Cpu64AdvCXXOmpOperator', 'Cpu64FsgCXXOperator', 'Cpu64FsgCXXOmpOperator'] + 'Cpu64AdvCXXOmpOperator', 'Cpu64FsgCXXOperator', 'Cpu64FsgCXXOmpOperator', + 'Cpu64NoopPetscOperator'] class Cpu64OperatorMixin: @@ -346,6 +347,15 @@ class Cpu64CXXNoopOmpOperator(Cpu64NoopOperator): LINEARIZE = True +class Cpu64NoopPetscOperator(Cpu64NoopOperator): + _Target = PetscTarget + + @classmethod + def _rcompile_wrapper(cls, **kwargs0): + kwargs0['language'] = 'petsc' + return super()._rcompile_wrapper(**kwargs0) + + class Cpu64AdvCOperator(Cpu64AdvOperator): _Target = CTarget @@ -355,6 +365,15 @@ class Cpu64AdvCXXOperator(Cpu64AdvOperator): LINEARIZE = True +class Cpu64AdvPetscOperator(Cpu64AdvOperator): + _Target = PetscTarget + + @classmethod + def _rcompile_wrapper(cls, **kwargs0): + kwargs0['language'] = 'petsc' + return super()._rcompile_wrapper(**kwargs0) + + class Cpu64AdvOmpOperator(Cpu64AdvOperator): _Target = COmpTarget diff --git a/devito/operator/registry.py b/devito/operator/registry.py index c8aac315b7..800cbe11a6 100644 --- a/devito/operator/registry.py +++ b/devito/operator/registry.py @@ -27,7 +27,7 @@ class OperatorRegistry(OrderedDict, metaclass=Singleton): _modes = ('noop', 'advanced', 'advanced-fsg') _languages = ('C', 'CXX', 'openmp', 'Copenmp', 'CXXopenmp', - 'openacc', 'cuda', 'hip', 'sycl') + 'openacc', 'cuda', 'hip', 'sycl', 'petsc') _accepted = _modes + tuple(product(_modes, _languages)) def add(self, operator, platform, mode, language='C'): diff --git a/devito/passes/iet/languages/C.py b/devito/passes/iet/languages/C.py index bfb0935a20..32d4f7f94b 100644 --- a/devito/passes/iet/languages/C.py +++ b/devito/passes/iet/languages/C.py @@ -6,6 +6,7 @@ from devito.passes.iet.orchestration import Orchestrator from devito.passes.iet.langbase import LangBB from devito.symbolics.extended_dtypes import c_complex, c_double_complex +from devito.petsc.utils import petsc_type_mappings __all__ = ['CBB', 'CDataManager', 'COrchestrator'] @@ -55,3 +56,9 @@ class CPrinter(BasePrinter, C99CodePrinter): def _print_ImaginaryUnit(self, expr): return '_Complex_I' + + +class PetscCPrinter(CPrinter): + _restrict_keyword = '' + + type_mappings = {**CPrinter.type_mappings, **petsc_type_mappings} diff --git a/devito/passes/iet/languages/targets.py b/devito/passes/iet/languages/targets.py index 3ca64e1c10..33bc4020ef 100644 --- a/devito/passes/iet/languages/targets.py +++ b/devito/passes/iet/languages/targets.py @@ -1,4 +1,5 @@ -from devito.passes.iet.languages.C import CDataManager, COrchestrator, CPrinter +from devito.passes.iet.languages.C import (CDataManager, COrchestrator, CPrinter, + PetscCPrinter) from devito.passes.iet.languages.CXX import CXXPrinter from devito.passes.iet.languages.openmp import (SimdOmpizer, Ompizer, DeviceOmpizer, OmpDataManager, DeviceOmpDataManager, @@ -8,7 +9,7 @@ from devito.passes.iet.instrument import instrument __all__ = ['CTarget', 'OmpTarget', 'COmpTarget', 'DeviceOmpTarget', 'DeviceAccTarget', - 'CXXTarget', 'CXXOmpTarget', 'DeviceCXXOmpTarget'] + 'CXXTarget', 'CXXOmpTarget', 'DeviceCXXOmpTarget', 'PetscTarget'] class Target: @@ -47,6 +48,10 @@ class COmpTarget(Target): Printer = CPrinter +class PetscTarget(CTarget): + Printer = PetscCPrinter + + OmpTarget = COmpTarget diff --git a/devito/petsc/iet/passes.py b/devito/petsc/iet/passes.py index 710f25a611..b31cfe21f1 100644 --- a/devito/petsc/iet/passes.py +++ b/devito/petsc/iet/passes.py @@ -17,7 +17,7 @@ SubMatrixStruct, Initialize, Finalize, ArgvSymbol) from devito.petsc.types.macros import petsc_func_begin_user from devito.petsc.iet.nodes import PetscMetaData -from devito.petsc.utils import core_metadata +from devito.petsc.utils import core_metadata, petsc_languages from devito.petsc.iet.routines import (CBBuilder, CCBBuilder, BaseObjectBuilder, CoupledObjectBuilder, BaseSetup, CoupledSetup, Solver, CoupledSolver, TimeDependent, @@ -34,6 +34,12 @@ def lower_petsc(iet, **kwargs): if not injectsolve_mapper: return iet, {} + if kwargs['language'] not in petsc_languages: + raise ValueError( + f"Expected 'language' to be one of " + f"{petsc_languages}, but got '{kwargs['language']}'" + ) + metadata = core_metadata() data = FindNodes(PetscMetaData).visit(iet) diff --git a/devito/petsc/initialize.py b/devito/petsc/initialize.py index 9126414658..80e3c7520c 100644 --- a/devito/petsc/initialize.py +++ b/devito/petsc/initialize.py @@ -3,7 +3,7 @@ from ctypes import POINTER, cast, c_char import atexit -from devito import Operator +from devito import Operator, switchconfig from devito.types import Symbol from devito.types.equation import PetscEq from devito.petsc.types import Initialize, Finalize @@ -20,14 +20,15 @@ def PetscInitialize(): # to generate these "dummy_ops" instead of using the Operator class. # This would prevent circular imports when initializing during import # from the PETSc module. - op_init = Operator( - [PetscEq(dummy, Initialize(dummy))], - name='kernel_init', opt='noop' - ) - op_finalize = Operator( - [PetscEq(dummy, Finalize(dummy))], - name='kernel_finalize', opt='noop' - ) + with switchconfig(language='petsc'): + op_init = Operator( + [PetscEq(dummy, Initialize(dummy))], + name='kernel_init', opt='noop' + ) + op_finalize = Operator( + [PetscEq(dummy, Finalize(dummy))], + name='kernel_finalize', opt='noop' + ) # `argv_bytes` must be a list so the memory address persists # `os.fsencode` should be preferred over `string().encode('utf-8')` diff --git a/devito/petsc/types/array.py b/devito/petsc/types/array.py index 38ac3bb9f3..e731e2cb53 100644 --- a/devito/petsc/types/array.py +++ b/devito/petsc/types/array.py @@ -106,10 +106,6 @@ def shape_allocated(self): @cached_property def _C_ctype(self): - # TODO: Switch to using PetscScalar instead of float/double - # TODO: Use cat $PETSC_DIR/$PETSC_ARCH/lib/petsc/conf/petscvariables - # | grep -E "PETSC_(SCALAR|PRECISION)" to determine the precision of - # the user's PETSc configuration. return POINTER(dtype_to_ctype(self.dtype)) @property diff --git a/devito/petsc/types/object.py b/devito/petsc/types/object.py index bf4339377d..9d2ee7b4fb 100644 --- a/devito/petsc/types/object.py +++ b/devito/petsc/types/object.py @@ -1,5 +1,5 @@ from ctypes import POINTER, c_char -from devito.tools import CustomDtype, dtype_to_cstr, as_tuple, CustomIntType +from devito.tools import CustomDtype, dtype_to_ctype, as_tuple, CustomIntType from devito.types import (LocalObject, LocalCompositeObject, ModuloDimension, TimeDimension, ArrayObject, CustomDimension) from devito.symbolics import Byref, cast @@ -188,7 +188,7 @@ class VecScatter(LocalObject): class StartPtr(LocalObject): def __init__(self, name, dtype): super().__init__(name=name) - self.dtype = CustomDtype(dtype_to_cstr(dtype), modifier=' *') + self.dtype = POINTER(dtype_to_ctype(dtype)) class SingleIS(LocalObject): diff --git a/devito/petsc/types/types.py b/devito/petsc/types/types.py index c7c3397181..782c5b7112 100644 --- a/devito/petsc/types/types.py +++ b/devito/petsc/types/types.py @@ -1,6 +1,8 @@ import sympy from devito.tools import Reconstructable, sympy_mutex +from devito.tools.dtypes_lowering import dtype_mapper +from devito.petsc.utils import petsc_variables class MetaData(sympy.Function, Reconstructable): @@ -132,6 +134,14 @@ class FieldData: def __init__(self, target=None, matvecs=None, formfuncs=None, formrhs=None, arrays=None, **kwargs): self._target = kwargs.get('target', target) + + petsc_precision = dtype_mapper[petsc_variables['PETSC_PRECISION']] + if self._target.dtype != petsc_precision: + raise TypeError( + f"Your target dtype must match the precision of your " + f"PETSc configuration. " + f"Expected {petsc_precision}, but got {self._target.dtype}." + ) self._matvecs = matvecs self._formfuncs = formfuncs self._formrhs = formrhs diff --git a/devito/petsc/utils.py b/devito/petsc/utils.py index 3dbb2c0d1a..07e8aa3157 100644 --- a/devito/petsc/utils.py +++ b/devito/petsc/utils.py @@ -1,8 +1,14 @@ import os +import ctypes +from pathlib import Path from devito.tools import memoized_func +class PetscOSError(OSError): + pass + + solver_mapper = { 'gmres': 'KSPGMRES', 'jacobi': 'PCJACOBI', @@ -12,42 +18,70 @@ @memoized_func def get_petsc_dir(): - # *** First try: via commonly used environment variables - for i in ['PETSC_DIR']: - petsc_dir = os.environ.get(i) - if petsc_dir: - return petsc_dir - # TODO: Raise error if PETSC_DIR is not set - return None + petsc_dir = os.environ.get('PETSC_DIR') + if petsc_dir is None: + raise PetscOSError("PETSC_DIR environment variable not set") + else: + petsc_dir = (Path(petsc_dir),) + petsc_arch = os.environ.get('PETSC_ARCH') + if petsc_arch is not None: + petsc_dir += (petsc_dir[0] / petsc_arch,) -@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 + petsc_installed = petsc_dir[-1] / 'include' / 'petscconf.h' + if not petsc_installed.is_file(): + raise PetscOSError("PETSc is not installed") + + return petsc_dir +@memoized_func def core_metadata(): petsc_dir = get_petsc_dir() - petsc_arch = get_petsc_arch() - - # Include directories - global_include = os.path.join(petsc_dir, 'include') - config_specific_include = os.path.join(petsc_dir, f'{petsc_arch}', 'include') - include_dirs = (global_include, config_specific_include) - # Lib directories - lib_dir = os.path.join(petsc_dir, f'{petsc_arch}', 'lib') + petsc_include = tuple([arch / 'include' for arch in petsc_dir]) + petsc_lib = tuple([arch / 'lib' for arch in petsc_dir]) return { 'includes': ('petscsnes.h', 'petscdmda.h'), - 'include_dirs': include_dirs, + 'include_dirs': petsc_include, 'libs': ('petsc'), - 'lib_dirs': lib_dir, - 'ldflags': ('-Wl,-rpath,%s' % lib_dir) + 'lib_dirs': petsc_lib, + 'ldflags': tuple([f"-Wl,-rpath,{lib}" for lib in petsc_lib]) } + + +@memoized_func +def get_petsc_variables(): + """ + Taken from https://www.firedrakeproject.org/_modules/firedrake/petsc.html + Get a dict of PETSc environment variables from the file: + $PETSC_DIR/$PETSC_ARCH/lib/petsc/conf/petscvariables + """ + try: + petsc_dir = get_petsc_dir() + except PetscOSError: + petsc_variables = {} + else: + path = [petsc_dir[-1], 'lib', 'petsc', 'conf', 'petscvariables'] + variables_path = Path(*path) + + with open(variables_path) as fh: + # Split lines on first '=' (assignment) + splitlines = (line.split("=", maxsplit=1) for line in fh.readlines()) + petsc_variables = {k.strip(): v.strip() for k, v in splitlines} + + return petsc_variables + + +petsc_variables = get_petsc_variables() + +# TODO: Check to see whether Petsc is compiled with +# 32-bit or 64-bit integers +# TODO: Check whether PetscScalar is a float or double +# and only map the right one +petsc_type_mappings = {ctypes.c_int: 'PetscInt', + ctypes.c_float: 'PetscScalar', + ctypes.c_double: 'PetscScalar'} + +petsc_languages = ['petsc'] diff --git a/examples/petsc/petsc_test.py b/examples/petsc/petsc_test.py index 69b63c99ec..5d93669d5f 100644 --- a/examples/petsc/petsc_test.py +++ b/examples/petsc/petsc_test.py @@ -1,7 +1,8 @@ import os import numpy as np -from devito import (Grid, Function, Eq, Operator, configuration) +from devito import (Grid, Function, Eq, Operator, configuration, + switchconfig) from devito.petsc import PETScSolve from devito.petsc.initialize import PetscInitialize configuration['compiler'] = 'custom' @@ -9,7 +10,6 @@ PetscInitialize() - nx = 81 ny = 81 @@ -24,6 +24,8 @@ petsc = PETScSolve([eq], u) -op = Operator(petsc) +with switchconfig(language='petsc'): + op = Operator(petsc) + op.apply() -op.apply() +print(op.ccode) diff --git a/tests/test_petsc.py b/tests/test_petsc.py index 1f0f8bc91e..d828371cb0 100644 --- a/tests/test_petsc.py +++ b/tests/test_petsc.py @@ -3,9 +3,9 @@ import pytest from conftest import skipif -from devito import (Grid, Function, TimeFunction, Eq, Operator, switchconfig, - configuration, norm) -from devito.ir.iet import (Call, ElementalFunction, Definition, DummyExpr, +from devito import (Grid, Function, TimeFunction, Eq, Operator, + configuration, norm, switchconfig) +from devito.ir.iet import (Call, ElementalFunction, FindNodes, retrieve_iteration_tree) from devito.types import Constant, LocalCompositeObject from devito.passes.iet.languages.C import CDataManager @@ -54,33 +54,6 @@ def test_petsc_local_object(): assert 'KSPConvergedReason reason;' in str(iet) -@skipif('petsc') -def test_petsc_functions(): - """ - Test C++ support for PETScArrays. - """ - grid = Grid((2, 2)) - x, y = grid.dimensions - - f0 = Function(name='f', grid=grid, space_order=2, dtype=np.float32) - f1 = Function(name='f', grid=grid, space_order=2, dtype=np.float64) - - ptr0 = PETScArray(name='ptr0', target=f0) - ptr1 = PETScArray(name='ptr1', target=f0, is_const=True) - ptr2 = PETScArray(name='ptr2', target=f1, is_const=True) - - defn0 = Definition(ptr0) - defn1 = Definition(ptr1) - defn2 = Definition(ptr2) - - expr = DummyExpr(ptr0.indexed[x, y], ptr1.indexed[x, y] + 1) - - assert str(defn0) == 'float *restrict ptr0_vec;' - assert str(defn1) == 'const float *restrict ptr1_vec;' - assert str(defn2) == 'const double *restrict ptr2_vec;' - assert str(expr) == 'ptr0[x][y] = ptr1[x][y] + 1;' - - @skipif('petsc') def test_petsc_subs(): """ @@ -112,7 +85,7 @@ def test_petsc_solve(): """ Test PETScSolve. """ - grid = Grid((2, 2)) + grid = Grid((2, 2), dtype=np.float64) f = Function(name='f', grid=grid, space_order=2) g = Function(name='g', grid=grid, space_order=2) @@ -121,7 +94,7 @@ def test_petsc_solve(): petsc = PETScSolve(eqn, f) - with switchconfig(openmp=False): + with switchconfig(language='petsc'): op = Operator(petsc, opt='noop') callable_roots = [meta_call.root for meta_call in op._func_table.values()] @@ -161,7 +134,7 @@ def test_multiple_petsc_solves(): """ Test multiple PETScSolves. """ - grid = Grid((2, 2)) + grid = Grid((2, 2), dtype=np.float64) f1 = Function(name='f1', grid=grid, space_order=2) g1 = Function(name='g1', grid=grid, space_order=2) @@ -175,7 +148,7 @@ def test_multiple_petsc_solves(): petsc1 = PETScSolve(eqn1, f1) petsc2 = PETScSolve(eqn2, f2) - with switchconfig(openmp=False): + with switchconfig(language='petsc'): op = Operator(petsc1+petsc2, opt='noop') callable_roots = [meta_call.root for meta_call in op._func_table.values()] @@ -189,9 +162,9 @@ def test_petsc_cast(): """ Test casting of PETScArray. """ - grid1 = Grid((2)) - grid2 = Grid((2, 2)) - grid3 = Grid((4, 5, 6)) + grid1 = Grid((2), dtype=np.float64) + grid2 = Grid((2, 2), dtype=np.float64) + grid3 = Grid((4, 5, 6), dtype=np.float64) f1 = Function(name='f1', grid=grid1, space_order=2) f2 = Function(name='f2', grid=grid2, space_order=4) @@ -205,37 +178,30 @@ def test_petsc_cast(): 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') - - cb1 = [meta_call.root for meta_call in op1._func_table.values()] - cb2 = [meta_call.root for meta_call in op2._func_table.values()] - cb3 = [meta_call.root for meta_call in op3._func_table.values()] + with switchconfig(language='petsc'): + op1 = Operator(petsc1) + op2 = Operator(petsc2) + op3 = Operator(petsc3) - assert 'float (*restrict x_f1) = ' + \ - '(float (*)) x_f1_vec;' in str(cb1[0]) - assert 'float (*restrict x_f2)[info.gxm] = ' + \ - '(float (*)[info.gxm]) x_f2_vec;' in str(cb2[0]) - assert 'float (*restrict x_f3)[info.gym][info.gxm] = ' + \ - '(float (*)[info.gym][info.gxm]) x_f3_vec;' in str(cb3[0]) + assert 'PetscScalar (* x_f1) = ' + \ + '(PetscScalar (*)) x_f1_vec;' in str(op1.ccode) + assert 'PetscScalar (* x_f2)[info.gxm] = ' + \ + '(PetscScalar (*)[info.gxm]) x_f2_vec;' in str(op2.ccode) + assert 'PetscScalar (* x_f3)[info.gym][info.gxm] = ' + \ + '(PetscScalar (*)[info.gym][info.gxm]) x_f3_vec;' in str(op3.ccode) @skipif('petsc') def test_LinearSolveExpr(): - grid = Grid((2, 2)) + grid = Grid((2, 2), dtype=np.float64) 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) - - # TODO: maybe expand this test now to check the fielddata etc - linsolveexpr = LinearSolveExpr(eqn.rhs) + linsolveexpr = LinearSolveExpr(eqn.rhs, fielddata=FieldData(target=f)) # Check the solver parameters assert linsolveexpr.solver_parameters == \ @@ -246,9 +212,9 @@ def test_LinearSolveExpr(): @skipif('petsc') def test_dmda_create(): - grid1 = Grid((2)) - grid2 = Grid((2, 2)) - grid3 = Grid((4, 5, 6)) + grid1 = Grid((2), dtype=np.float64) + grid2 = Grid((2, 2), dtype=np.float64) + grid3 = Grid((4, 5, 6), dtype=np.float64) f1 = Function(name='f1', grid=grid1, space_order=2) f2 = Function(name='f2', grid=grid2, space_order=4) @@ -262,7 +228,7 @@ def test_dmda_create(): petsc2 = PETScSolve(eqn2, f2) petsc3 = PETScSolve(eqn3, f3) - with switchconfig(openmp=False): + with switchconfig(language='petsc'): op1 = Operator(petsc1, opt='noop') op2 = Operator(petsc2, opt='noop') op3 = Operator(petsc3, opt='noop') @@ -282,13 +248,14 @@ def test_dmda_create(): @skipif('petsc') def test_cinterface_petsc_struct(): - grid = Grid(shape=(11, 11)) + grid = Grid(shape=(11, 11), dtype=np.float64) f = Function(name='f', grid=grid, space_order=2) eq = Eq(f.laplace, 10) petsc = PETScSolve(eq, f) name = "foo" - with switchconfig(openmp=False): + + with switchconfig(language='petsc'): op = Operator(petsc, name=name) # Trigger the generation of a .c and a .h files @@ -561,7 +528,7 @@ def test_callback_arguments(): """ Test the arguments of each callback function. """ - grid = Grid((2, 2)) + grid = Grid((2, 2), dtype=np.float64) f1 = Function(name='f1', grid=grid, space_order=2) g1 = Function(name='g1', grid=grid, space_order=2) @@ -570,7 +537,7 @@ def test_callback_arguments(): petsc1 = PETScSolve(eqn1, f1) - with switchconfig(openmp=False): + with switchconfig(language='petsc'): op = Operator(petsc1) mv = op._func_table['MatMult0'].root @@ -586,7 +553,7 @@ def test_callback_arguments(): @skipif('petsc') def test_petsc_struct(): - grid = Grid((2, 2)) + grid = Grid((2, 2), dtype=np.float64) f1 = Function(name='f1', grid=grid, space_order=2) g1 = Function(name='g1', grid=grid, space_order=2) @@ -599,7 +566,7 @@ def test_petsc_struct(): eqn2 = Eq(f1, g1*mu2) - with switchconfig(openmp=False): + with switchconfig(language='petsc'): op = Operator([eqn2] + petsc1) arguments = op.arguments() @@ -621,19 +588,20 @@ def test_apply(): 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) + pn = Function(name='pn', grid=grid, space_order=2) + rhs = Function(name='rhs', grid=grid, space_order=2) mu = Constant(name='mu', value=2.0) eqn = Eq(pn.laplace*mu, rhs, subdomain=grid.interior) petsc = PETScSolve(eqn, pn) - # Build the op - op = Operator(petsc) + with switchconfig(language='petsc'): + # Build the op + op = Operator(petsc) - # Check the Operator runs without errors - op.apply() + # Check the Operator runs without errors + op.apply() # Verify that users can override `mu` mu_new = Constant(name='mu_new', value=4.0) @@ -643,7 +611,7 @@ def test_apply(): @skipif('petsc') def test_petsc_frees(): - grid = Grid((2, 2)) + grid = Grid((2, 2), dtype=np.float64) f = Function(name='f', grid=grid, space_order=2) g = Function(name='g', grid=grid, space_order=2) @@ -651,7 +619,7 @@ def test_petsc_frees(): eqn = Eq(f.laplace, g) petsc = PETScSolve(eqn, f) - with switchconfig(openmp=False): + with switchconfig(language='petsc'): op = Operator(petsc) frees = op.body.frees @@ -667,7 +635,7 @@ def test_petsc_frees(): @skipif('petsc') def test_calls_to_callbacks(): - grid = Grid((2, 2)) + grid = Grid((2, 2), dtype=np.float64) f = Function(name='f', grid=grid, space_order=2) g = Function(name='g', grid=grid, space_order=2) @@ -675,7 +643,7 @@ def test_calls_to_callbacks(): eqn = Eq(f.laplace, g) petsc = PETScSolve(eqn, f) - with switchconfig(openmp=False): + with switchconfig(language='petsc'): op = Operator(petsc) ccode = str(op.ccode) @@ -694,27 +662,28 @@ def test_start_ptr(): 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) + grid = Grid((11, 11), dtype=np.float64) + u1 = TimeFunction(name='u1', grid=grid, space_order=2) eq1 = Eq(u1.dt, u1.laplace, subdomain=grid.interior) petsc1 = PETScSolve(eq1, u1.forward) - with switchconfig(openmp=False): + with switchconfig(language='petsc'): op1 = Operator(petsc1) # Verify the case with modulo time stepping - assert 'float * u1_ptr0 = t1*localsize0 + (float*)(u1_vec->data);' in str(op1) + assert ('PetscScalar * u1_ptr0 = t1*localsize0 + ' + '(PetscScalar*)(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) + u2 = TimeFunction(name='u2', grid=grid, space_order=2, save=5) eq2 = Eq(u2.dt, u2.laplace, subdomain=grid.interior) petsc2 = PETScSolve(eq2, u2.forward) - with switchconfig(openmp=False): + with switchconfig(language='petsc'): op2 = Operator(petsc2) - assert 'float * u2_ptr0 = (time + 1)*localsize0 + ' + \ - '(float*)(u2_vec->data);' in str(op2) + assert ('PetscScalar * u2_ptr0 = (time + 1)*localsize0 + ' + '(PetscScalar*)(u2_vec->data);') in str(op2) @skipif('petsc') @@ -726,15 +695,17 @@ def test_time_loop(): - Only assign/update the modulo dimensions required by any of the PETSc callback functions. """ - grid = Grid((11, 11)) + grid = Grid((11, 11), dtype=np.float64) # Modulo time stepping u1 = TimeFunction(name='u1', grid=grid, space_order=2) v1 = Function(name='v1', grid=grid, space_order=2) eq1 = Eq(v1.laplace, u1) petsc1 = PETScSolve(eq1, v1) - with switchconfig(openmp=False): + + with switchconfig(language='petsc'): op1 = Operator(petsc1) + op1.apply(time_M=3) body1 = str(op1.body) rhs1 = str(op1._func_table['FormRHS0'].root.ccode) @@ -748,8 +719,10 @@ def test_time_loop(): v2 = Function(name='v2', grid=grid, space_order=2, save=5) eq2 = Eq(v2.laplace, u2) petsc2 = PETScSolve(eq2, v2) - with switchconfig(openmp=False): + + with switchconfig(language='petsc'): op2 = Operator(petsc2) + op2.apply(time_M=3) body2 = str(op2.body) rhs2 = str(op2._func_table['FormRHS0'].root.ccode) @@ -760,8 +733,10 @@ def test_time_loop(): # used in one of the callback functions eq3 = Eq(v1.laplace, u1 + u1.forward) petsc3 = PETScSolve(eq3, v1) - with switchconfig(openmp=False): + + with switchconfig(language='petsc'): op3 = Operator(petsc3) + op3.apply(time_M=3) body3 = str(op3.body) rhs3 = str(op3._func_table['FormRHS0'].root.ccode) @@ -776,14 +751,40 @@ def test_time_loop(): petsc4 = PETScSolve(eq4, v1) eq5 = Eq(v2.laplace, u1) petsc5 = PETScSolve(eq5, v2) - with switchconfig(openmp=False): + + with switchconfig(language='petsc'): op4 = Operator(petsc4 + petsc5) + op4.apply(time_M=3) body4 = str(op4.body) assert 'ctx0.t0 = t0' in body4 assert body4.count('ctx0.t0 = t0') == 1 +@skipif('petsc') +def test_solve_output(): + """ + Verify that PETScSolve returns the correct output for + simple cases e.g with the identity matrix. + """ + grid = Grid(shape=(11, 11), dtype=np.float64) + + u = Function(name='u', grid=grid, space_order=2) + v = Function(name='v', grid=grid, space_order=2) + + # Solving Ax=b where A is the identity matrix + v.data[:] = 5.0 + eqn = Eq(u, v) + petsc = PETScSolve(eqn, target=u) + + with switchconfig(language='petsc'): + op = Operator(petsc) + # Check the solve function returns the correct output + op.apply() + + assert np.allclose(u.data, v.data) + + class TestCoupledLinear: # The coupled interface can be used even for uncoupled problems, meaning # the equations will be solved within a single matrix system. @@ -796,7 +797,7 @@ class TestCoupledLinear: def test_coupled_vs_non_coupled(self): grid = Grid(shape=(11, 11), dtype=np.float64) - functions = [Function(name=n, grid=grid, space_order=2, dtype=np.float64) + functions = [Function(name=n, grid=grid, space_order=2) for n in ['e', 'f', 'g', 'h']] e, f, g, h = functions @@ -810,9 +811,9 @@ def test_coupled_vs_non_coupled(self): petsc1 = PETScSolve(eq1, target=e) petsc2 = PETScSolve(eq2, target=g) - with switchconfig(openmp=False): + with switchconfig(language='petsc'): op1 = Operator(petsc1 + petsc2, opt='noop') - op1.apply() + op1.apply() enorm1 = norm(e) gnorm1 = norm(g) @@ -825,9 +826,10 @@ def test_coupled_vs_non_coupled(self): # TODO: Need more friendly API for coupled - just # using a dict for now petsc3 = PETScSolve({e: [eq1], g: [eq2]}) - with switchconfig(openmp=False): + + with switchconfig(language='petsc'): op2 = Operator(petsc3, opt='noop') - op2.apply() + op2.apply() enorm2 = norm(e) gnorm2 = norm(g) @@ -856,9 +858,9 @@ def test_coupled_vs_non_coupled(self): @skipif('petsc') def test_coupled_structs(self): - grid = Grid(shape=(11, 11)) + grid = Grid(shape=(11, 11), dtype=np.float64) - functions = [Function(name=n, grid=grid, space_order=2, dtype=np.float64) + functions = [Function(name=n, grid=grid, space_order=2) for n in ['e', 'f', 'g', 'h']] e, f, g, h = functions @@ -868,7 +870,8 @@ def test_coupled_structs(self): petsc = PETScSolve({f: [eq1], h: [eq2]}) name = "foo" - with switchconfig(openmp=False): + + with switchconfig(language='petsc'): op = Operator(petsc, name=name) # Trigger the generation of a .c and a .h files @@ -899,7 +902,7 @@ def test_coupled_structs(self): def test_coupled_frees(self): grid = Grid(shape=(11, 11), dtype=np.float64) - functions = [Function(name=n, grid=grid, space_order=2, dtype=np.float64) + functions = [Function(name=n, grid=grid, space_order=2) for n in ['e', 'f', 'g', 'h']] e, f, g, h = functions @@ -910,7 +913,7 @@ def test_coupled_frees(self): petsc1 = PETScSolve({e: [eq1], f: [eq2]}) petsc2 = PETScSolve({e: [eq1], f: [eq2], g: [eq3]}) - with switchconfig(openmp=False): + with switchconfig(language='petsc'): op1 = Operator(petsc1, opt='noop') op2 = Operator(petsc2, opt='noop') @@ -941,9 +944,9 @@ def test_coupled_frees(self): @skipif('petsc') def test_dmda_dofs(self): - grid = Grid(shape=(11, 11)) + grid = Grid(shape=(11, 11), dtype=np.float64) - functions = [Function(name=n, grid=grid, space_order=2, dtype=np.float64) + functions = [Function(name=n, grid=grid, space_order=2) for n in ['e', 'f', 'g', 'h']] e, f, g, h = functions @@ -955,7 +958,7 @@ def test_dmda_dofs(self): petsc2 = PETScSolve({e: [eq1], f: [eq2]}) petsc3 = PETScSolve({e: [eq1], f: [eq2], g: [eq3]}) - with switchconfig(openmp=False): + with switchconfig(language='petsc'): op1 = Operator(petsc1, opt='noop') op2 = Operator(petsc2, opt='noop') op3 = Operator(petsc3, opt='noop') @@ -975,9 +978,9 @@ def test_dmda_dofs(self): @skipif('petsc') def test_submatrices(self): - grid = Grid(shape=(11, 11)) + grid = Grid(shape=(11, 11), dtype=np.float64) - functions = [Function(name=n, grid=grid, space_order=2, dtype=np.float64) + functions = [Function(name=n, grid=grid, space_order=2) for n in ['e', 'f', 'g', 'h']] e, f, g, h = functions From 6cdb7d01c1ce9dd962e6f430c42e372793bdc0c6 Mon Sep 17 00:00:00 2001 From: ZoeLeibowitz <90093761+ZoeLeibowitz@users.noreply.github.com> Date: Mon, 31 Mar 2025 17:19:10 +0100 Subject: [PATCH 08/84] compiler: Temp fix for memory leaks. (#70) - Temp fix until VecReplaceArray works properly in conjunction with PetscMemoryAllocator --- devito/mpi/distributed.py | 2 +- devito/petsc/iet/passes.py | 23 +++++++------ devito/petsc/iet/routines.py | 64 +++++++++++++++++++++++------------- devito/petsc/solve.py | 22 ++++++++----- devito/petsc/types/object.py | 20 ++--------- devito/petsc/utils.py | 27 +++++++++++---- tests/test_petsc.py | 15 +++++---- 7 files changed, 98 insertions(+), 75 deletions(-) diff --git a/devito/mpi/distributed.py b/devito/mpi/distributed.py index 1fc587367e..b9bd568fd0 100644 --- a/devito/mpi/distributed.py +++ b/devito/mpi/distributed.py @@ -91,7 +91,7 @@ def devito_mpi_finalize(): """ Finalize MPI, if initialized by Devito. """ - global init_by_devito + global init_by_devito # noqa if init_by_devito and MPI.Is_initialized() and not MPI.Is_finalized(): MPI.Finalize() diff --git a/devito/petsc/iet/passes.py b/devito/petsc/iet/passes.py index b31cfe21f1..3c73bfd391 100644 --- a/devito/petsc/iet/passes.py +++ b/devito/petsc/iet/passes.py @@ -11,7 +11,7 @@ from devito.types.basic import DataSymbol from devito.tools import frozendict from devito.petsc.types import (PetscMPIInt, PetscErrorCode, MultipleFieldData, - PointerIS, Mat, LocalVec, GlobalVec, CallbackMat, SNES, + PointerIS, Mat, CallbackVec, Vec, CallbackMat, SNES, DummyArg, PetscInt, PointerDM, PointerMat, MatReuse, CallbackPointerIS, CallbackPointerDM, JacobianStruct, SubMatrixStruct, Initialize, Finalize, ArgvSymbol) @@ -40,14 +40,13 @@ def lower_petsc(iet, **kwargs): f"{petsc_languages}, but got '{kwargs['language']}'" ) - metadata = core_metadata() data = FindNodes(PetscMetaData).visit(iet) if any(filter(lambda i: isinstance(i.expr.rhs, Initialize), data)): - return initialize(iet), metadata + return initialize(iet), core_metadata() if any(filter(lambda i: isinstance(i.expr.rhs, Finalize), data)): - return finalize(iet), metadata + return finalize(iet), core_metadata() unique_grids = {i.expr.rhs.grid for (i,) in injectsolve_mapper.values()} # Assumption is that all solves are on the same grid @@ -79,7 +78,7 @@ def lower_petsc(iet, **kwargs): body = core + tuple(setup) + (BlankLine,) + iet.body.body body = iet.body._rebuild(body=body) iet = iet._rebuild(body=body) - metadata.update({'efuncs': tuple(efuncs.values())}) + metadata = {**core_metadata(), 'efuncs': tuple(efuncs.values())} return iet, metadata @@ -218,13 +217,13 @@ def populate_matrix_context(efuncs, objs): 'rowidx': PetscInt('rowidx'), 'colidx': PetscInt('colidx'), 'J': Mat('J'), - 'X': GlobalVec('X'), - 'xloc': LocalVec('xloc'), - 'Y': GlobalVec('Y'), - 'yloc': LocalVec('yloc'), - 'F': GlobalVec('F'), - 'floc': LocalVec('floc'), - 'B': GlobalVec('B'), + 'X': Vec('X'), + 'xloc': CallbackVec('xloc'), + 'Y': Vec('Y'), + 'yloc': CallbackVec('yloc'), + 'F': Vec('F'), + 'floc': CallbackVec('floc'), + 'B': Vec('B'), 'nfields': PetscInt('nfields'), 'irow': PointerIS(name='irow'), 'icol': PointerIS(name='icol'), diff --git a/devito/petsc/iet/routines.py b/devito/petsc/iet/routines.py index 4c28231cec..bd3c3305b0 100644 --- a/devito/petsc/iet/routines.py +++ b/devito/petsc/iet/routines.py @@ -16,7 +16,7 @@ MatShellSetOp, PetscMetaData) from devito.petsc.iet.utils import petsc_call, petsc_struct from devito.petsc.utils import solver_mapper -from devito.petsc.types import (DM, Mat, LocalVec, GlobalVec, KSP, PC, SNES, +from devito.petsc.types import (DM, Mat, CallbackVec, Vec, KSP, PC, SNES, PetscInt, StartPtr, PointerIS, PointerDM, VecScatter, DMCast, JacobianStructCast, JacobianStruct, SubMatrixStruct, CallbackDM) @@ -448,8 +448,13 @@ def _create_form_rhs_body(self, body, fielddata): 'VecRestoreArray', [sobjs['blocal'], Byref(b_arr._C_symbol)] ) + dm_restore_local_bvec = petsc_call( + 'DMRestoreLocalVector', [dmda, Byref(sobjs['blocal'])] + ) + body = body._rebuild(body=body.body + ( - dm_local_to_global_begin, dm_local_to_global_end, vec_restore_array + dm_local_to_global_begin, dm_local_to_global_end, vec_restore_array, + dm_restore_local_bvec )) stacks = ( @@ -870,10 +875,10 @@ def _build(self): targets = self.fielddata.targets base_dict = { 'Jac': Mat(sreg.make_name(prefix='J')), - 'xglobal': GlobalVec(sreg.make_name(prefix='xglobal')), - 'xlocal': LocalVec(sreg.make_name(prefix='xlocal')), - 'bglobal': GlobalVec(sreg.make_name(prefix='bglobal')), - 'blocal': LocalVec(sreg.make_name(prefix='blocal')), + 'xglobal': Vec(sreg.make_name(prefix='xglobal')), + 'xlocal': Vec(sreg.make_name(prefix='xlocal')), + 'bglobal': Vec(sreg.make_name(prefix='bglobal')), + 'blocal': CallbackVec(sreg.make_name(prefix='blocal')), 'ksp': KSP(sreg.make_name(prefix='ksp')), 'pc': PC(sreg.make_name(prefix='pc')), 'snes': SNES(sreg.make_name(prefix='snes')), @@ -939,9 +944,9 @@ def _extend_build(self, base_dict): name=f'{key}ctx', fields=objs['subctx'].fields, ) - base_dict[f'{key}X'] = LocalVec(f'{key}X') - base_dict[f'{key}Y'] = LocalVec(f'{key}Y') - base_dict[f'{key}F'] = LocalVec(f'{key}F') + base_dict[f'{key}X'] = CallbackVec(f'{key}X') + base_dict[f'{key}Y'] = CallbackVec(f'{key}Y') + base_dict[f'{key}F'] = CallbackVec(f'{key}F') return base_dict @@ -953,22 +958,22 @@ def _target_dependent(self, base_dict): base_dict[f'{name}_ptr'] = StartPtr( sreg.make_name(prefix=f'{name}_ptr'), t.dtype ) - base_dict[f'xlocal{name}'] = LocalVec( + base_dict[f'xlocal{name}'] = CallbackVec( sreg.make_name(prefix=f'xlocal{name}'), liveness='eager' ) - base_dict[f'Fglobal{name}'] = LocalVec( + base_dict[f'Fglobal{name}'] = CallbackVec( sreg.make_name(prefix=f'Fglobal{name}'), liveness='eager' ) - base_dict[f'Xglobal{name}'] = LocalVec( + base_dict[f'Xglobal{name}'] = CallbackVec( sreg.make_name(prefix=f'Xglobal{name}') ) - base_dict[f'xglobal{name}'] = GlobalVec( + base_dict[f'xglobal{name}'] = Vec( sreg.make_name(prefix=f'xglobal{name}') ) - base_dict[f'blocal{name}'] = LocalVec( + base_dict[f'blocal{name}'] = CallbackVec( sreg.make_name(prefix=f'blocal{name}'), liveness='eager' ) - base_dict[f'bglobal{name}'] = GlobalVec( + base_dict[f'bglobal{name}'] = Vec( sreg.make_name(prefix=f'bglobal{name}') ) base_dict[f'da{name}'] = DM( @@ -1021,6 +1026,12 @@ def _setup(self): global_x = petsc_call('DMCreateGlobalVector', [dmda, Byref(sobjs['xglobal'])]) + local_x = petsc_call('DMCreateLocalVector', + [dmda, Byref(sobjs['xlocal'])]) + + get_local_size = petsc_call('VecGetSize', + [sobjs['xlocal'], Byref(sobjs['localsize'])]) + global_b = petsc_call('DMCreateGlobalVector', [dmda, Byref(sobjs['bglobal'])]) @@ -1078,6 +1089,8 @@ def _setup(self): snes_set_jac, snes_set_type, global_x, + local_x, + get_local_size, global_b, snes_get_ksp, ksp_set_tols, @@ -1235,9 +1248,6 @@ def _execute_solve(self): rhs_call = petsc_call(rhs_callback.name, [sobjs['dmda'], sobjs['bglobal']]) - local_x = petsc_call('DMCreateLocalVector', - [dmda, Byref(sobjs['xlocal'])]) - vec_replace_array = self.timedep.replace_array(target) dm_local_to_global_x = petsc_call( @@ -1253,13 +1263,15 @@ def _execute_solve(self): dmda, sobjs['xglobal'], insert_vals, sobjs['xlocal']] ) + vec_reset_array = self.timedep.reset_array(target) + run_solver_calls = (struct_assignment,) + ( rhs_call, - local_x ) + vec_replace_array + ( dm_local_to_global_x, snes_solve, dm_global_to_local_x, + vec_reset_array, BlankLine, ) return List(body=run_solver_calls) @@ -1402,7 +1414,16 @@ def replace_array(self, target): target.function._C_field_data, target.function._C_symbol ) xlocal = sobjs.get(f'xlocal{target.name}', sobjs['xlocal']) - return (petsc_call('VecReplaceArray', [xlocal, field_from_ptr]),) + return (petsc_call('VecPlaceArray', [xlocal, field_from_ptr]),) + + def reset_array(self, target): + """ + """ + sobjs = self.sobjs + xlocal = sobjs.get(f'xlocal{target.name}', sobjs['xlocal']) + return ( + petsc_call('VecResetArray', [xlocal]) + ) def assign_time_iters(self, struct): return [] @@ -1530,7 +1551,6 @@ def replace_array(self, target): caster = cast(target.dtype, '*') return ( - petsc_call('VecGetSize', [xlocal, Byref(sobjs['localsize'])]), DummyExpr( start_ptr, caster( @@ -1538,7 +1558,7 @@ def replace_array(self, target): ) + Mul(target_time, sobjs['localsize']), init=True ), - petsc_call('VecReplaceArray', [xlocal, start_ptr]) + petsc_call('VecPlaceArray', [xlocal, start_ptr]) ) return super().replace_array(target) diff --git a/devito/petsc/solve.py b/devito/petsc/solve.py index 971fc8678b..334df4c802 100644 --- a/devito/petsc/solve.py +++ b/devito/petsc/solve.py @@ -70,29 +70,33 @@ def build_function_eqns(self, eq, target, arrays): formfunc = self.make_formfunc(eq, F_target, arrays, targets) formrhs = self.make_rhs(eq, b, arrays) - return tuple(expr.subs(self.time_mapper) for expr in (formfunc, formrhs)) + return (formfunc, formrhs) def build_matvec_eqns(self, eq, target, arrays): b, F_target, targets = separate_eqn(eq, target) if not F_target: return None matvec = self.make_matvec(eq, F_target, arrays, targets) - return matvec.subs(self.time_mapper) + return matvec def make_matvec(self, eq, F_target, arrays, targets): - rhs = arrays['x'] if isinstance(eq, EssentialBC) else F_target.subs( - targets_to_arrays(arrays['x'], targets) - ) + if isinstance(eq, EssentialBC): + rhs = arrays['x'] + else: + rhs = F_target.subs(targets_to_arrays(arrays['x'], targets)) + rhs = rhs.subs(self.time_mapper) return Eq(arrays['y'], rhs, subdomain=eq.subdomain) def make_formfunc(self, eq, F_target, arrays, targets): - rhs = 0. if isinstance(eq, EssentialBC) else F_target.subs( - targets_to_arrays(arrays['x'], targets) - ) + if isinstance(eq, EssentialBC): + rhs = 0. + else: + rhs = F_target.subs(targets_to_arrays(arrays['x'], targets)) + rhs = rhs.subs(self.time_mapper) return Eq(arrays['f'], rhs, subdomain=eq.subdomain) def make_rhs(self, eq, b, arrays): - rhs = 0. if isinstance(eq, EssentialBC) else b + rhs = 0. if isinstance(eq, EssentialBC) else b.subs(self.time_mapper) return Eq(arrays['b'], rhs, subdomain=eq.subdomain) def generate_arrays(self, target): diff --git a/devito/petsc/types/object.py b/devito/petsc/types/object.py index 9d2ee7b4fb..fcbf06ea69 100644 --- a/devito/petsc/types/object.py +++ b/devito/petsc/types/object.py @@ -66,28 +66,14 @@ def _C_free_priority(self): return 2 -class LocalVec(LocalObject): +class CallbackVec(LocalObject): """ - PETSc local vector object (Vec). - A local vector has ghost locations that contain values that are - owned by other MPI ranks. + PETSc vector object (Vec). """ dtype = CustomDtype('Vec') -class CallbackGlobalVec(LocalVec): - """ - PETSc global vector object (Vec). For example, used for coupled - solves inside the `WholeFormFunc` callback. - """ - - -class GlobalVec(LocalVec): - """ - PETSc global vector object (Vec). - A global vector is a parallel vector that has no duplicate values - between MPI ranks. A global vector has no ghost locations. - """ +class Vec(CallbackVec): @property def _C_free(self): return petsc_call('VecDestroy', [Byref(self.function)]) diff --git a/devito/petsc/utils.py b/devito/petsc/utils.py index 07e8aa3157..3d491703e7 100644 --- a/devito/petsc/utils.py +++ b/devito/petsc/utils.py @@ -76,12 +76,25 @@ def get_petsc_variables(): petsc_variables = get_petsc_variables() -# TODO: Check to see whether Petsc is compiled with -# 32-bit or 64-bit integers -# TODO: Check whether PetscScalar is a float or double -# and only map the right one -petsc_type_mappings = {ctypes.c_int: 'PetscInt', - ctypes.c_float: 'PetscScalar', - ctypes.c_double: 'PetscScalar'} + +def get_petsc_type_mappings(): + try: + petsc_precision = petsc_variables['PETSC_PRECISION'] + except KeyError: + mapper = {} + else: + petsc_scalar = 'PetscScalar' + # TODO: Check to see whether Petsc is compiled with + # 32-bit or 64-bit integers + mapper = {ctypes.c_int: 'PetscInt'} + + if petsc_precision == 'single': + mapper[ctypes.c_float] = petsc_scalar + elif petsc_precision == 'double': + mapper[ctypes.c_double] = petsc_scalar + return mapper + + +petsc_type_mappings = get_petsc_type_mappings() petsc_languages = ['petsc'] diff --git a/tests/test_petsc.py b/tests/test_petsc.py index d828371cb0..47d6faabcd 100644 --- a/tests/test_petsc.py +++ b/tests/test_petsc.py @@ -9,7 +9,7 @@ FindNodes, retrieve_iteration_tree) from devito.types import Constant, LocalCompositeObject from devito.passes.iet.languages.C import CDataManager -from devito.petsc.types import (DM, Mat, LocalVec, PetscMPIInt, KSP, +from devito.petsc.types import (DM, Mat, Vec, PetscMPIInt, KSP, PC, KSPConvergedReason, PETScArray, LinearSolveExpr, FieldData, MultipleFieldData) from devito.petsc.solve import PETScSolve, separate_eqn, centre_stencil @@ -33,7 +33,7 @@ def test_petsc_local_object(): """ lo0 = DM('da', stencil_width=1) lo1 = Mat('A') - lo2 = LocalVec('x') + lo2 = Vec('x') lo3 = PetscMPIInt('size') lo4 = KSP('ksp') lo5 = PC('pc') @@ -590,7 +590,7 @@ def test_apply(): pn = Function(name='pn', grid=grid, space_order=2) rhs = Function(name='rhs', grid=grid, space_order=2) - mu = Constant(name='mu', value=2.0) + mu = Constant(name='mu', value=2.0, dtype=np.float64) eqn = Eq(pn.laplace*mu, rhs, subdomain=grid.interior) @@ -604,7 +604,7 @@ def test_apply(): op.apply() # Verify that users can override `mu` - mu_new = Constant(name='mu_new', value=4.0) + mu_new = Constant(name='mu_new', value=4.0, dtype=np.float64) op.apply(mu=mu_new) @@ -627,9 +627,10 @@ def test_petsc_frees(): # Check the frees appear in the following order assert str(frees[0]) == 'PetscCall(VecDestroy(&(bglobal0)));' assert str(frees[1]) == 'PetscCall(VecDestroy(&(xglobal0)));' - assert str(frees[2]) == 'PetscCall(MatDestroy(&(J0)));' - assert str(frees[3]) == 'PetscCall(SNESDestroy(&(snes0)));' - assert str(frees[4]) == 'PetscCall(DMDestroy(&(da0)));' + assert str(frees[2]) == 'PetscCall(VecDestroy(&(xlocal0)));' + assert str(frees[3]) == 'PetscCall(MatDestroy(&(J0)));' + assert str(frees[4]) == 'PetscCall(SNESDestroy(&(snes0)));' + assert str(frees[5]) == 'PetscCall(DMDestroy(&(da0)));' @skipif('petsc') From 9d389da3c8cb5732656a0fbeeef132d734edf0f7 Mon Sep 17 00:00:00 2001 From: ZoeLeibowitz <90093761+ZoeLeibowitz@users.noreply.github.com> Date: Mon, 31 Mar 2025 20:18:50 +0100 Subject: [PATCH 09/84] workflows: Push PETSc/Devito docker image to Docker Hub (#71) --- .github/workflows/pytest-petsc.yml | 11 +++++++++++ devito/mpi/distributed.py | 2 +- 2 files changed, 12 insertions(+), 1 deletion(-) diff --git a/.github/workflows/pytest-petsc.yml b/.github/workflows/pytest-petsc.yml index 090748b222..ce74eceb4a 100644 --- a/.github/workflows/pytest-petsc.yml +++ b/.github/workflows/pytest-petsc.yml @@ -46,9 +46,20 @@ jobs: - name: Checkout devito uses: actions/checkout@v4 + - name: Log in to DockerHub + uses: docker/login-action@v2 + with: + username: ${{ secrets.DOCKER_USERNAME }} + password: ${{ secrets.DOCKER_PASSWORD }} + - 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 + docker tag devito_img ${{ secrets.DOCKER_USERNAME }}/devito_img:latest + + - name: Push Docker image to DockerHub + run: | + docker push ${{ secrets.DOCKER_USERNAME }}/devito_img:latest - name: Set run prefix run: | diff --git a/devito/mpi/distributed.py b/devito/mpi/distributed.py index b9bd568fd0..63e3643fe0 100644 --- a/devito/mpi/distributed.py +++ b/devito/mpi/distributed.py @@ -91,7 +91,7 @@ def devito_mpi_finalize(): """ Finalize MPI, if initialized by Devito. """ - global init_by_devito # noqa + global init_by_devito # noqa: F824 if init_by_devito and MPI.Is_initialized() and not MPI.Is_finalized(): MPI.Finalize() From fb4b96fc701fb7c7cf23ed81fe49ace50fdf90e8 Mon Sep 17 00:00:00 2001 From: ZoeLeibowitz <90093761+ZoeLeibowitz@users.noreply.github.com> Date: Wed, 2 Apr 2025 08:54:39 +0100 Subject: [PATCH 10/84] Compiler: Add basic initguess callback (#72) * Compiler: Add basic initguess callback - to be improved --- devito/petsc/iet/routines.py | 92 +++++++++++++++++++ devito/petsc/solve.py | 23 +++++ devito/petsc/types/types.py | 7 +- .../seismic/tutorials/13_LSRTM_acoustic.ipynb | 5 +- tests/test_petsc.py | 73 ++++++++++++++- 5 files changed, 195 insertions(+), 5 deletions(-) diff --git a/devito/petsc/iet/routines.py b/devito/petsc/iet/routines.py index bd3c3305b0..82a0798246 100644 --- a/devito/petsc/iet/routines.py +++ b/devito/petsc/iet/routines.py @@ -47,6 +47,7 @@ def __init__(self, **kwargs): self._matvecs = [] self._formfuncs = [] self._formrhs = [] + self._initialguesses = [] self._make_core() self._efuncs = self._uxreplace_efuncs() @@ -88,6 +89,10 @@ def formfuncs(self): def formrhs(self): return self._formrhs + @property + def initialguesses(self): + return self._initialguesses + @property def user_struct_callback(self): return self._user_struct_callback @@ -97,6 +102,8 @@ def _make_core(self): self._make_matvec(fielddata, fielddata.matvecs) self._make_formfunc(fielddata) self._make_formrhs(fielddata) + if fielddata.initialguess: + self._make_initialguess(fielddata) self._make_user_struct_callback() def _make_matvec(self, fielddata, matvecs, prefix='MatMult'): @@ -483,6 +490,84 @@ def _create_form_rhs_body(self, body, fielddata): return Uxreplace(subs).visit(formrhs_body) + def _make_initialguess(self, fielddata): + initguess = fielddata.initialguess + sobjs = self.solver_objs + + # Compile initital guess `eqns` into an IET via recursive compilation + irs, _ = self.rcompile( + initguess, options={'mpi': False}, sregistry=self.sregistry, + concretize_mapper=self.concretize_mapper + ) + body_init_guess = self._create_initial_guess_body( + List(body=irs.uiet.body), fielddata + ) + objs = self.objs + cb = PETScCallable( + self.sregistry.make_name(prefix='FormInitialGuess'), + body_init_guess, + retval=objs['err'], + parameters=(sobjs['callbackdm'], objs['xloc']) + ) + self._initialguesses.append(cb) + self._efuncs[cb.name] = cb + + def _create_initial_guess_body(self, body, fielddata): + linsolve_expr = self.injectsolve.expr.rhs + objs = self.objs + sobjs = self.solver_objs + + dmda = sobjs['callbackdm'] + ctx = objs['dummyctx'] + + x_arr = fielddata.arrays['x'] + + vec_get_array = petsc_call( + 'VecGetArray', [objs['xloc'], Byref(x_arr._C_symbol)] + ) + + dm_get_local_info = petsc_call( + 'DMDAGetLocalInfo', [dmda, Byref(linsolve_expr.localinfo)] + ) + + body = self.timedep.uxreplace_time(body) + + fields = self._dummy_fields(body) + self._struct_params.extend(fields) + + dm_get_app_context = petsc_call( + 'DMGetApplicationContext', [dmda, Byref(ctx._C_symbol)] + ) + + vec_restore_array = petsc_call( + 'VecRestoreArray', [objs['xloc'], Byref(x_arr._C_symbol)] + ) + + body = body._rebuild(body=body.body + (vec_restore_array,)) + + stacks = ( + vec_get_array, + dm_get_app_context, + dm_get_local_info + ) + + # Dereference function data in struct + dereference_funcs = [Dereference(i, ctx) for i in + fields if isinstance(i.function, AbstractFunction)] + + body = CallableBody( + List(body=[body]), + init=(objs['begin_user'],), + stacks=stacks+tuple(dereference_funcs), + retstmt=(Call('PetscFunctionReturn', arguments=[0]),) + ) + + # Replace non-function data with pointer to data in struct + subs = {i._C_symbol: FieldFromPointer(i._C_symbol, ctx) for + i in fields if not isinstance(i.function, AbstractFunction)} + + return Uxreplace(subs).visit(body) + def _make_user_struct_callback(self): """ This is the struct initialised inside the main kernel and @@ -1250,6 +1335,12 @@ def _execute_solve(self): vec_replace_array = self.timedep.replace_array(target) + if self.cbbuilder.initialguesses: + initguess = self.cbbuilder.initialguesses[0] + initguess_call = petsc_call(initguess.name, [dmda, sobjs['xlocal']]) + else: + initguess_call = None + dm_local_to_global_x = petsc_call( 'DMLocalToGlobal', [dmda, sobjs['xlocal'], insert_vals, sobjs['xglobal']] @@ -1268,6 +1359,7 @@ def _execute_solve(self): run_solver_calls = (struct_assignment,) + ( rhs_call, ) + vec_replace_array + ( + initguess_call, dm_local_to_global_x, snes_solve, dm_global_to_local_x, diff --git a/devito/petsc/solve.py b/devito/petsc/solve.py index 334df4c802..a4917dfe8f 100644 --- a/devito/petsc/solve.py +++ b/devito/petsc/solve.py @@ -57,11 +57,18 @@ def generate_field_data(self, eqns, target, arrays): ) matvecs = [self.build_matvec_eqns(eq, target, arrays) for eq in eqns] + initialguess = [ + eq for eq in + (self.make_initial_guess(e, target, arrays) for e in eqns) + if eq is not None + ] + return FieldData( target=target, matvecs=matvecs, formfuncs=formfuncs, formrhs=formrhs, + initialguess=initialguess, arrays=arrays ) @@ -99,6 +106,22 @@ def make_rhs(self, eq, b, arrays): rhs = 0. if isinstance(eq, EssentialBC) else b.subs(self.time_mapper) return Eq(arrays['b'], rhs, subdomain=eq.subdomain) + def make_initial_guess(self, eq, target, arrays): + """ + Enforce initial guess to satisfy essential BCs. + # TODO: For time-stepping, only enforce these once outside the time loop + and use the previous time-step solution as the initial guess for next time step. + # TODO: Extend this to "coupled". + """ + if isinstance(eq, EssentialBC): + assert eq.lhs == target + return Eq( + arrays['x'], eq.rhs, + subdomain=eq.subdomain + ) + else: + return None + def generate_arrays(self, target): return { p: PETScArray(name=f'{p}_{target.name}', diff --git a/devito/petsc/types/types.py b/devito/petsc/types/types.py index 782c5b7112..bfbdd4972b 100644 --- a/devito/petsc/types/types.py +++ b/devito/petsc/types/types.py @@ -132,7 +132,7 @@ def eval(cls, *args): class FieldData: def __init__(self, target=None, matvecs=None, formfuncs=None, formrhs=None, - arrays=None, **kwargs): + initialguess=None, arrays=None, **kwargs): self._target = kwargs.get('target', target) petsc_precision = dtype_mapper[petsc_variables['PETSC_PRECISION']] @@ -145,6 +145,7 @@ def __init__(self, target=None, matvecs=None, formfuncs=None, formrhs=None, self._matvecs = matvecs self._formfuncs = formfuncs self._formrhs = formrhs + self._initialguess = initialguess self._arrays = arrays @property @@ -163,6 +164,10 @@ def formfuncs(self): def formrhs(self): return self._formrhs + @property + def initialguess(self): + return self._initialguess + @property def arrays(self): return self._arrays diff --git a/examples/seismic/tutorials/13_LSRTM_acoustic.ipynb b/examples/seismic/tutorials/13_LSRTM_acoustic.ipynb index 71bcab0f1a..42a266dd47 100644 --- a/examples/seismic/tutorials/13_LSRTM_acoustic.ipynb +++ b/examples/seismic/tutorials/13_LSRTM_acoustic.ipynb @@ -304,7 +304,7 @@ }, { "cell_type": "code", - "execution_count": 6, + "execution_count": null, "metadata": {}, "outputs": [], "source": [ @@ -329,12 +329,13 @@ " dm_true = (solver.model.vp.data**(-2) - model0.vp.data**(-2))\n", " \n", " objective = 0.\n", + " u0 = None\n", " for i in range(nshots):\n", " \n", " #Observed Data using Born's operator\n", " geometry.src_positions[0, :] = source_locations[i, :]\n", "\n", - " _, u0, _ = solver.forward(vp=model0.vp, save=True)\n", + " _, u0, _ = solver.forward(vp=model0.vp, save=True, u=u0)\n", " \n", " _, _, _,_ = solver.jacobian(dm_true, vp=model0.vp, rec = d_obs)\n", " \n", diff --git a/tests/test_petsc.py b/tests/test_petsc.py index 47d6faabcd..e89f4c4476 100644 --- a/tests/test_petsc.py +++ b/tests/test_petsc.py @@ -4,7 +4,7 @@ from conftest import skipif from devito import (Grid, Function, TimeFunction, Eq, Operator, - configuration, norm, switchconfig) + configuration, norm, switchconfig, SubDomain) from devito.ir.iet import (Call, ElementalFunction, FindNodes, retrieve_iteration_tree) from devito.types import Constant, LocalCompositeObject @@ -12,7 +12,8 @@ from devito.petsc.types import (DM, Mat, Vec, PetscMPIInt, KSP, PC, KSPConvergedReason, PETScArray, LinearSolveExpr, FieldData, MultipleFieldData) -from devito.petsc.solve import PETScSolve, separate_eqn, centre_stencil +from devito.petsc.solve import (PETScSolve, separate_eqn, centre_stencil, + EssentialBC) from devito.petsc.iet.nodes import Expression from devito.petsc.initialize import PetscInitialize @@ -786,6 +787,74 @@ def test_solve_output(): assert np.allclose(u.data, v.data) +@skipif('petsc') +def test_essential_bcs(): + """ + Verify that PETScSolve returns the correct output with + essential boundary conditions. + """ + class SubTop(SubDomain): + name = 'subtop' + + def define(self, dimensions): + x, y = dimensions + return {x: x, y: ('right', 1)} + sub1 = SubTop() + + class SubBottom(SubDomain): + name = 'subbottom' + + def define(self, dimensions): + x, y = dimensions + return {x: x, y: ('left', 1)} + sub2 = SubBottom() + + class SubLeft(SubDomain): + name = 'subleft' + + def define(self, dimensions): + x, y = dimensions + return {x: ('left', 1), y: y} + sub3 = SubLeft() + + class SubRight(SubDomain): + name = 'subright' + + def define(self, dimensions): + x, y = dimensions + return {x: ('right', 1), y: y} + sub4 = SubRight() + + subdomains = (sub1, sub2, sub3, sub4) + grid = Grid(shape=(11, 11), subdomains=subdomains, dtype=np.float64) + + u = Function(name='u', grid=grid, space_order=2) + v = Function(name='v', grid=grid, space_order=2) + + # Solving Ax=b where A is the identity matrix + v.data[:] = 5.0 + eqn = Eq(u, v) + + bcs = [EssentialBC(u, 1., subdomain=sub1)] # top + bcs += [EssentialBC(u, 2., subdomain=sub2)] # bottom + bcs += [EssentialBC(u, 3., subdomain=sub3)] # left + bcs += [EssentialBC(u, 4., subdomain=sub4)] # right + + petsc = PETScSolve([eqn]+bcs, target=u) + + with switchconfig(language='petsc'): + op = Operator(petsc) + op.apply() + + # Check u is equal to v on the interior + assert np.allclose(u.data[1:-1, 1:-1], v.data[1:-1, 1:-1]) + # Check u satisfies the boundary conditions + assert np.allclose(u.data[1:-1, -1], 1.0) # top + assert np.allclose(u.data[1:-1, 0], 2.0) # bottom + assert np.allclose(u.data[0, 1:-1], 3.0) # left + assert np.allclose(u.data[-1, 1:-1], 4.0) # right + + class TestCoupledLinear: # The coupled interface can be used even for uncoupled problems, meaning # the equations will be solved within a single matrix system. From 33b289632c26342c403d8469b6e266a76bfece06 Mon Sep 17 00:00:00 2001 From: ZoeLeibowitz <90093761+ZoeLeibowitz@users.noreply.github.com> Date: Wed, 2 Apr 2025 12:31:17 +0100 Subject: [PATCH 11/84] Staggered grids (#73) --- .github/workflows/pytest-petsc.yml | 4 + devito/petsc/iet/routines.py | 8 +- devito/petsc/types/array.py | 7 +- examples/petsc/seismic/staggered_acoustic.py | 97 ++++++++++++++++++++ 4 files changed, 112 insertions(+), 4 deletions(-) create mode 100644 examples/petsc/seismic/staggered_acoustic.py diff --git a/.github/workflows/pytest-petsc.yml b/.github/workflows/pytest-petsc.yml index ce74eceb4a..bc3ed58876 100644 --- a/.github/workflows/pytest-petsc.yml +++ b/.github/workflows/pytest-petsc.yml @@ -79,6 +79,10 @@ jobs: run: | ${{ env.RUN_CMD }} mpiexec -n 1 pytest --cov --cov-config=.coveragerc --cov-report=xml ${{ env.TESTS }} + - name: Test examples + run: | + ${{ env.RUN_CMD }} mpiexec -n 1 python3 --cov --cov-config=.coveragerc --cov-report=xml examples/petsc/seismic/staggered_acoustic.py + - name: Upload coverage to Codecov if: "!contains(matrix.name, 'docker')" uses: codecov/codecov-action@v4 diff --git a/devito/petsc/iet/routines.py b/devito/petsc/iet/routines.py index 82a0798246..5a000e57a6 100644 --- a/devito/petsc/iet/routines.py +++ b/devito/petsc/iet/routines.py @@ -1373,9 +1373,11 @@ def spatial_body(self): spatial_body = [] # TODO: remove the iters[0] for tree in retrieve_iteration_tree(self.iters[0]): - root = filter_iterations(tree, key=lambda i: i.dim.is_Space)[0] - if self.injectsolve in FindNodes(PetscMetaData).visit(root): - spatial_body.append(root) + root = filter_iterations(tree, key=lambda i: i.dim.is_Space) + if root: + root = root[0] + if self.injectsolve in FindNodes(PetscMetaData).visit(root): + spatial_body.append(root) spatial_body, = spatial_body return spatial_body diff --git a/devito/petsc/types/array.py b/devito/petsc/types/array.py index e731e2cb53..b07b73db48 100644 --- a/devito/petsc/types/array.py +++ b/devito/petsc/types/array.py @@ -59,7 +59,8 @@ def __dtype_setup__(cls, **kwargs): @classmethod def __indices_setup__(cls, *args, **kwargs): - dimensions = kwargs['target'].space_dimensions + target = kwargs['target'] + dimensions = tuple(target.indices[d] for d in target.space_dimensions) if args: indices = args else: @@ -92,6 +93,10 @@ def shape(self): def space_order(self): return self.target.space_order + @property + def is_Staggered(self): + return self.target.staggered is not None + @property def localinfo(self): return self._localinfo diff --git a/examples/petsc/seismic/staggered_acoustic.py b/examples/petsc/seismic/staggered_acoustic.py new file mode 100644 index 0000000000..4f0c91ed49 --- /dev/null +++ b/examples/petsc/seismic/staggered_acoustic.py @@ -0,0 +1,97 @@ +from devito import * +import os +import numpy as np +from examples.seismic.source import DGaussSource, TimeAxis +from devito.petsc import PETScSolve +from devito.petsc.initialize import PetscInitialize +configuration['compiler'] = 'custom' +os.environ['CC'] = 'mpicc' + + +# PETSc implementation of devito/examples/seismic/tutorials/05_staggered_acoustic.ipynb +# Test staggered grid implementation with PETSc + +PetscInitialize() + +extent = (2000., 2000.) +shape = (81, 81) + +x = SpaceDimension( + name='x', spacing=Constant(name='h_x', value=extent[0]/(shape[0]-1), dtype=np.float64) +) +z = SpaceDimension( + name='z', spacing=Constant(name='h_z', value=extent[1]/(shape[1]-1), dtype=np.float64) +) + +grid = Grid(extent=extent, shape=shape, dimensions=(x, z), dtype=np.float64) + +# Timestep size +t0, tn = 0., 200. +dt = 1e2*(1. / np.sqrt(2.)) / 60. +time_range = TimeAxis(start=t0, stop=tn, step=dt) + +src = DGaussSource(name='src', grid=grid, f0=0.01, time_range=time_range, a=0.004) +src.coordinates.data[:] = [1000., 1000.] + +# Now we create the velocity and pressure fields +# NOTE/TODO: PETSc does not yet fully support VectorTimeFunctions. Ideally, +# it should use the new "coupled" machinery +p2 = TimeFunction(name='p2', grid=grid, staggered=NODE, space_order=2, time_order=1) +vx2 = TimeFunction(name='vx2', grid=grid, staggered=(x,), space_order=2, time_order=1) +vz2 = TimeFunction(name='vz2', grid=grid, staggered=(z,), space_order=2, time_order=1) + +t = grid.stepping_dim +time = grid.time_dim + +# We need some initial conditions +V_p = 4.0 +density = 1. + +ro = 1/density +l2m = V_p*V_p*density + +# The source injection term +src_p_2 = src.inject(field=p2.forward, expr=src) + +# 2nd order acoustic according to fdelmoc +v_x_2 = Eq(vx2.dt, ro * p2.dx) +v_z_2 = Eq(vz2.dt, ro * p2.dz) + +petsc_v_x_2 = PETScSolve(v_x_2, target=vx2.forward) +petsc_v_z_2 = PETScSolve(v_z_2, target=vz2.forward) + +p_2 = Eq(p2.dt, l2m * (vx2.forward.dx + vz2.forward.dz)) + +petsc_p_2 = PETScSolve(p_2, target=p2.forward, solver_parameters={'ksp_rtol': 1e-7}) + +with switchconfig(language='petsc'): + op_2 = Operator(petsc_v_x_2 + petsc_v_z_2 + petsc_p_2 + src_p_2, opt='noop') + op_2(time=src.time_range.num-1, dt=dt) + +norm_p2 = norm(p2) +assert np.isclose(norm_p2, .35098, atol=1e-4, rtol=0) + + +# 4th order acoustic according to fdelmoc +p4 = TimeFunction(name='p4', grid=grid, staggered=NODE, space_order=4) +vx4 = TimeFunction(name='vx4', grid=grid, staggered=(x,), space_order=4, time_order=1) +vz4 = TimeFunction(name='vz4', grid=grid, staggered=(z,), space_order=4, time_order=1) + +src_p_4 = src.inject(field=p4.forward, expr=src) + +v_x_4 = Eq(vx4.dt, ro * p4.dx) +v_z_4 = Eq(vz4.dt, ro * p4.dz) + +petsc_v_x_4 = PETScSolve(v_x_4, target=vx4.forward) +petsc_v_z_4 = PETScSolve(v_z_4, target=vz4.forward) + +p_4 = Eq(p4.dt, l2m * (vx4.forward.dx + vz4.forward.dz)) + +petsc_p_4 = PETScSolve(p_4, target=p4.forward, solver_parameters={'ksp_rtol': 1e-7}) + +with switchconfig(language='petsc'): + op_4 = Operator(petsc_v_x_4 + petsc_v_z_4 + petsc_p_4 + src_p_4, opt='noop') + op_4(time=src.time_range.num-1, dt=dt) + +norm_p4 = norm(p4) +assert np.isclose(norm_p4, .33736, atol=1e-4, rtol=0) From 515046a5c368b0728de88ff5441764ba3dfa2f40 Mon Sep 17 00:00:00 2001 From: ZoeLeibowitz <90093761+ZoeLeibowitz@users.noreply.github.com> Date: Wed, 2 Apr 2025 14:09:48 +0100 Subject: [PATCH 12/84] Drop unnecessary IndexedData check (#74) --- devito/ir/equations/algorithms.py | 2 +- devito/passes/iet/definitions.py | 5 ++--- 2 files changed, 3 insertions(+), 4 deletions(-) diff --git a/devito/ir/equations/algorithms.py b/devito/ir/equations/algorithms.py index 5828879012..49f65a2c5d 100644 --- a/devito/ir/equations/algorithms.py +++ b/devito/ir/equations/algorithms.py @@ -191,7 +191,7 @@ def concretize_subdims(exprs, **kwargs): """ sregistry = kwargs.get('sregistry') - # Update based on changes in #2509 + # To be updated based on changes in #2509 mapper = kwargs.get('concretize_mapper', {}) rebuilt = {} # Rebuilt implicit dims etc which are shared between dimensions diff --git a/devito/passes/iet/definitions.py b/devito/passes/iet/definitions.py index 859de274b6..fcf3aafe3e 100644 --- a/devito/passes/iet/definitions.py +++ b/devito/passes/iet/definitions.py @@ -19,7 +19,7 @@ SizeOf, VOID, pow_to_mul) from devito.tools import as_mapper, as_list, as_tuple, filter_sorted, flatten from devito.types import (Array, ComponentAccess, CustomDimension, DeviceMap, - DeviceRM, Eq, Symbol, IndexedData) + DeviceRM, Eq, Symbol) __all__ = ['DataManager', 'DeviceAwareDataManager', 'Storage'] @@ -451,8 +451,7 @@ 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 - if isinstance(i.base, IndexedData)}, key=lambda i: i.name) + bases = sorted({i.base for i in indexeds}, 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 From f1d7f0f7cdefabe690ae3b65b53e9de1db3c6f1b Mon Sep 17 00:00:00 2001 From: ZoeLeibowitz <90093761+ZoeLeibowitz@users.noreply.github.com> Date: Wed, 2 Apr 2025 14:28:33 +0100 Subject: [PATCH 13/84] Separate PETSc install into single image (#75) --- .github/workflows/pytest-petsc.yml | 7 +++---- docker/Dockerfile.devito | 23 +---------------------- docker/Dockerfile.petsc | 25 +++++++++++++++++++++++++ 3 files changed, 29 insertions(+), 26 deletions(-) create mode 100644 docker/Dockerfile.petsc diff --git a/.github/workflows/pytest-petsc.yml b/.github/workflows/pytest-petsc.yml index bc3ed58876..234363c06e 100644 --- a/.github/workflows/pytest-petsc.yml +++ b/.github/workflows/pytest-petsc.yml @@ -54,16 +54,15 @@ jobs: - 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 - docker tag devito_img ${{ secrets.DOCKER_USERNAME }}/devito_img:latest + docker build -f docker/Dockerfile.devito --build-arg base=zoeleibowitz/petsc_image:latest --tag zoeleibowitz/petsc_devito_image:latest . - name: Push Docker image to DockerHub run: | - docker push ${{ secrets.DOCKER_USERNAME }}/devito_img:latest + docker push ${{ secrets.DOCKER_USERNAME }}/petsc_devito_image:latest - name: Set run prefix run: | - echo "RUN_CMD=docker run --rm -t -e CODECOV_TOKEN=${{ secrets.CODECOV_TOKEN }} --name testrun devito_img" >> $GITHUB_ENV + echo "RUN_CMD=docker run --rm -t -e CODECOV_TOKEN=${{ secrets.CODECOV_TOKEN }} --name testrun zoeleibowitz/petsc_devito_image:latest" >> $GITHUB_ENV id: set-run - name: Set tests diff --git a/docker/Dockerfile.devito b/docker/Dockerfile.devito index ba41309fd9..8e4148555e 100644 --- a/docker/Dockerfile.devito +++ b/docker/Dockerfile.devito @@ -4,26 +4,8 @@ # Base image with compilers ARG base=devitocodes/bases:cpu-gcc -ARG petscinstall="" -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 +FROM $base AS builder # User/Group Ids ARG USER_ID=1000 @@ -87,9 +69,6 @@ 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.petsc b/docker/Dockerfile.petsc new file mode 100644 index 0000000000..6bf0800a8e --- /dev/null +++ b/docker/Dockerfile.petsc @@ -0,0 +1,25 @@ +############################################################## +# Dockerfile.petsc: Installs PETSc +############################################################## + +# Base image with compilers +# TODO: to be updated, but made some additions to Dockerfile.cpu so need to +# use the one from my dockerhub +ARG base=zoeleibowitz/bases:cpu-gcc + +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 + +ENV PETSC_DIR="/opt/petsc/petsc" +ENV PETSC_ARCH="devito_build" From 42e0e70ba108c1807406cd35e8ae9d3c1d518e5d Mon Sep 17 00:00:00 2001 From: ZoeLeibowitz <90093761+ZoeLeibowitz@users.noreply.github.com> Date: Wed, 2 Apr 2025 19:49:53 +0100 Subject: [PATCH 14/84] Add some examples (#76) --- .github/workflows/pytest-petsc.yml | 6 +- devito/petsc/types/array.py | 21 +- examples/petsc/Poisson/01_poisson.py | 112 +++++++ examples/petsc/Poisson/02_laplace.py | 125 +++++++ examples/petsc/cfd/01_navierstokes.py | 307 ++++++++++++++++++ examples/petsc/random/01_helmholtz.py | 259 +++++++++++++++ ...d_acoustic.py => 01_staggered_acoustic.py} | 0 7 files changed, 817 insertions(+), 13 deletions(-) create mode 100644 examples/petsc/Poisson/01_poisson.py create mode 100644 examples/petsc/Poisson/02_laplace.py create mode 100644 examples/petsc/cfd/01_navierstokes.py create mode 100644 examples/petsc/random/01_helmholtz.py rename examples/petsc/seismic/{staggered_acoustic.py => 01_staggered_acoustic.py} (100%) diff --git a/.github/workflows/pytest-petsc.yml b/.github/workflows/pytest-petsc.yml index 234363c06e..0d44614b16 100644 --- a/.github/workflows/pytest-petsc.yml +++ b/.github/workflows/pytest-petsc.yml @@ -80,7 +80,11 @@ jobs: - name: Test examples run: | - ${{ env.RUN_CMD }} mpiexec -n 1 python3 --cov --cov-config=.coveragerc --cov-report=xml examples/petsc/seismic/staggered_acoustic.py + ${{ env.RUN_CMD }} mpiexec -n 1 python3 examples/petsc/seismic/01_staggered_acoustic.py + ${{ env.RUN_CMD }} mpiexec -n 1 python3 examples/petsc/cfd/01_navierstokes.py + ${{ env.RUN_CMD }} mpiexec -n 1 python3 examples/petsc/Poisson/01_poisson.py + ${{ env.RUN_CMD }} mpiexec -n 1 python3 examples/petsc/Poisson/02_laplace.py + ${{ env.RUN_CMD }} mpiexec -n 1 python3 examples/petsc/random/01_helmholtz.py - name: Upload coverage to Codecov if: "!contains(matrix.name, 'docker')" diff --git a/devito/petsc/types/array.py b/devito/petsc/types/array.py index b07b73db48..54bbb4a9f4 100644 --- a/devito/petsc/types/array.py +++ b/devito/petsc/types/array.py @@ -5,7 +5,6 @@ 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, as_tuple from devito.symbolics import FieldFromComposite @@ -27,9 +26,6 @@ class PETScArray(ArrayBasic, Differentiable): _data_alignment = False - # Default method for the finite difference approximation weights computation. - _default_fd = 'taylor' - __rkwargs__ = (AbstractFunction.__rkwargs__ + ('target', 'liveness', 'coefficients', 'localinfo')) @@ -38,15 +34,8 @@ def __init_finalize__(self, *args, **kwargs): self._target = kwargs.get('target') self._ndim = kwargs['ndim'] = len(self._target.space_dimensions) self._dimensions = kwargs['dimensions'] = self._target.space_dimensions - super().__init_finalize__(*args, **kwargs) - - # Symbolic (finite difference) coefficients - 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._coefficients = self._target.coefficients self._localinfo = kwargs.get('localinfo', None) @property @@ -93,6 +82,10 @@ def shape(self): def space_order(self): return self.target.space_order + @property + def staggered(self): + return self.target.staggered + @property def is_Staggered(self): return self.target.staggered is not None @@ -113,6 +106,10 @@ def shape_allocated(self): def _C_ctype(self): return POINTER(dtype_to_ctype(self.dtype)) + @cached_property + def _fd_priority(self): + return self.target._fd_priority + @property def symbolic_shape(self): field_from_composites = [ diff --git a/examples/petsc/Poisson/01_poisson.py b/examples/petsc/Poisson/01_poisson.py new file mode 100644 index 0000000000..6d89091cd6 --- /dev/null +++ b/examples/petsc/Poisson/01_poisson.py @@ -0,0 +1,112 @@ +import os +import numpy as np + +from devito import (Grid, Function, Eq, Operator, switchconfig, + configuration, SubDomain) + +from devito.petsc import PETScSolve, EssentialBC +from devito.petsc.initialize import PetscInitialize +configuration['compiler'] = 'custom' +os.environ['CC'] = 'mpicc' + +# Solving pn.laplace = 2x(y - 1)(y - 2x + xy + 2)e^(x-y) +# Constant zero Dirichlet BCs. + +PetscInitialize() + + +# Subdomains to implement BCs +class SubTop(SubDomain): + name = 'subtop' + + def define(self, dimensions): + x, y = dimensions + return {x: x, y: ('right', 1)} + + +class SubBottom(SubDomain): + name = 'subbottom' + + def define(self, dimensions): + x, y = dimensions + return {x: x, y: ('left', 1)} + + +class SubLeft(SubDomain): + name = 'subleft' + + def define(self, dimensions): + x, y = dimensions + return {x: ('left', 1), y: y} + + +class SubRight(SubDomain): + name = 'subright' + + def define(self, dimensions): + x, y = dimensions + return {x: ('right', 1), y: y} + + +sub1 = SubTop() +sub2 = SubBottom() +sub3 = SubLeft() +sub4 = SubRight() + +subdomains = (sub1, sub2, sub3, sub4) + + +def analytical(x, y): + return np.float64(np.exp(x-y) * x * (1-x) * y * (1-y)) + + +Lx = np.float64(1.) +Ly = np.float64(1.) + +n_values = list(range(13, 174, 10)) +dx = np.array([Lx/(n-1) for n in n_values]) +errors = [] + + +for n in n_values: + + grid = Grid( + shape=(n, n), extent=(Lx, Ly), subdomains=subdomains, dtype=np.float64 + ) + + phi = Function(name='phi', grid=grid, space_order=2, dtype=np.float64) + rhs = Function(name='rhs', grid=grid, space_order=2, dtype=np.float64) + + eqn = Eq(rhs, phi.laplace, subdomain=grid.interior) + + tmpx = np.linspace(0, Lx, n).astype(np.float64) + tmpy = np.linspace(0, Ly, n).astype(np.float64) + Y, X = np.meshgrid(tmpx, tmpy) + + rhs.data[:] = np.float64( + 2.0*X*(Y-1.0)*(Y - 2.0*X + X*Y + 2.0) + ) * np.float64(np.exp(X-Y)) + + # # Create boundary condition expressions using subdomains + bcs = [EssentialBC(phi, np.float64(0.), subdomain=sub1)] + bcs += [EssentialBC(phi, np.float64(0.), subdomain=sub2)] + bcs += [EssentialBC(phi, np.float64(0.), subdomain=sub3)] + bcs += [EssentialBC(phi, np.float64(0.), subdomain=sub4)] + + exprs = [eqn] + bcs + petsc = PETScSolve(exprs, target=phi, solver_parameters={'ksp_rtol': 1e-8}) + + with switchconfig(language='petsc'): + op = Operator(petsc) + op.apply() + + phi_analytical = analytical(X, Y) + + diff = phi_analytical[1:-1, 1:-1] - phi.data[1:-1, 1:-1] + error = np.linalg.norm(diff) / np.linalg.norm(phi_analytical[1:-1, 1:-1]) + errors.append(error) + +slope, _ = np.polyfit(np.log(dx), np.log(errors), 1) + +assert slope > 1.9 +assert slope < 2.1 diff --git a/examples/petsc/Poisson/02_laplace.py b/examples/petsc/Poisson/02_laplace.py new file mode 100644 index 0000000000..780da10ec5 --- /dev/null +++ b/examples/petsc/Poisson/02_laplace.py @@ -0,0 +1,125 @@ +import os +import numpy as np + +from devito import (Grid, Function, Eq, Operator, SubDomain, + configuration, switchconfig) +from devito.petsc import PETScSolve, EssentialBC +from devito.petsc.initialize import PetscInitialize +configuration['compiler'] = 'custom' +os.environ['CC'] = 'mpicc' + +PetscInitialize() + +# Laplace equation, solving phi.laplace = 0 + +# Constant Dirichlet BCs: +# phi(x, 0) = 0 +# phi(0, y) = 0 +# phi(1, y) = 0 +# phi(x, 1) = f(x) = sin(pi*x) + +# The analytical solution is: +# phi(x, y) = sinh(pi*y)*sin(pi*x)/sinh(pi) + + +# Subdomains to implement BCs +class SubTop(SubDomain): + name = 'subtop' + + def define(self, dimensions): + x, y = dimensions + return {x: x, y: ('right', 1)} + + +class SubBottom(SubDomain): + name = 'subbottom' + + def define(self, dimensions): + x, y = dimensions + return {x: x, y: ('left', 1)} + + +class SubLeft(SubDomain): + name = 'subleft' + + def define(self, dimensions): + x, y = dimensions + return {x: ('left', 1), y: y} + + +class SubRight(SubDomain): + name = 'subright' + + def define(self, dimensions): + x, y = dimensions + return {x: ('right', 1), y: y} + + +sub1 = SubTop() +sub2 = SubBottom() +sub3 = SubLeft() +sub4 = SubRight() + +subdomains = (sub1, sub2, sub3, sub4) + + +Lx = np.float64(1.) +Ly = np.float64(1.) + + +def analytical(x, y, Lx, Ly): + tmp = np.float64(np.pi)/Lx + numerator = np.float64(np.sinh(tmp*y)) * np.float64(np.sin(tmp*x)) + return numerator / np.float64(np.sinh(tmp*Ly)) + + +n_values = list(range(13, 174, 10)) +dx = np.array([Lx/(n-1) for n in n_values]) +errors = [] + + +for n in n_values: + + grid = Grid( + shape=(n, n), extent=(Lx, Ly), subdomains=subdomains, dtype=np.float64 + ) + + phi = Function(name='phi', grid=grid, space_order=2, dtype=np.float64) + rhs = Function(name='rhs', grid=grid, space_order=2, dtype=np.float64) + + phi.data[:] = np.float64(0.0) + rhs.data[:] = np.float64(0.0) + + eqn = Eq(rhs, phi.laplace, subdomain=grid.interior) + + tmpx = np.linspace(0, Lx, n).astype(np.float64) + tmpy = np.linspace(0, Ly, n).astype(np.float64) + Y, X = np.meshgrid(tmpx, tmpy) + + # Create boundary condition expressions using subdomains + bc_func = Function(name='bcs', grid=grid, space_order=2, dtype=np.float64) + bc_func.data[:] = np.float64(0.0) + bc_func.data[:, -1] = np.float64(np.sin(tmpx*np.pi)) + + bcs = [EssentialBC(phi, bc_func, subdomain=sub1)] # top + bcs += [EssentialBC(phi, bc_func, subdomain=sub2)] # bottom + bcs += [EssentialBC(phi, bc_func, subdomain=sub3)] # left + bcs += [EssentialBC(phi, bc_func, subdomain=sub4)] # right + + exprs = [eqn] + bcs + petsc = PETScSolve(exprs, target=phi, solver_parameters={'ksp_rtol': 1e-8}) + + with switchconfig(language='petsc'): + op = Operator(petsc) + op.apply() + + phi_analytical = analytical(X, Y, Lx, Ly) + + diff = phi_analytical[1:-1, 1:-1] - phi.data[1:-1, 1:-1] + error = np.linalg.norm(diff) / np.linalg.norm(phi_analytical[1:-1, 1:-1]) + errors.append(error) + +slope, _ = np.polyfit(np.log(dx), np.log(errors), 1) + +assert slope > 1.9 +assert slope < 2.1 diff --git a/examples/petsc/cfd/01_navierstokes.py b/examples/petsc/cfd/01_navierstokes.py new file mode 100644 index 0000000000..01e96e1afe --- /dev/null +++ b/examples/petsc/cfd/01_navierstokes.py @@ -0,0 +1,307 @@ +import os +import numpy as np + +from devito import (Grid, TimeFunction, Function, Constant, Eq, + Operator, norm, SubDomain, switchconfig, configuration) +from devito.symbolics import retrieve_functions, INT + +from devito.petsc import PETScSolve, EssentialBC +from devito.petsc.initialize import PetscInitialize +configuration['compiler'] = 'custom' +os.environ['CC'] = 'mpicc' + + +PetscInitialize() + +# Chorin's projection method +# Explicit time-stepping + +# Physical parameters +rho = Constant(name='rho', dtype=np.float64) +nu = Constant(name='nu', dtype=np.float64) + +rho.data = np.float64(1.) +nu.data = np.float64(1./10.) + +Lx = 1. +Ly = Lx + +# Number of grid points in each direction +nx = 41 +ny = 41 + +# mesh spacing +dx = Lx/(nx-1) +dy = Ly/(ny-1) +so = 2 + + +# Use subdomains just for pressure field for now +class SubTop(SubDomain): + name = 'subtop' + + def __init__(self, S_O): + super().__init__() + self.S_O = S_O + + def define(self, dimensions): + x, y = dimensions + return {x: ('middle', 1, 1), y: ('right', self.S_O//2)} + + +class SubBottom(SubDomain): + name = 'subbottom' + + def __init__(self, S_O): + super().__init__() + self.S_O = S_O + + def define(self, dimensions): + x, y = dimensions + return {x: ('middle', 1, 1), y: ('left', self.S_O//2)} + + +class SubLeft(SubDomain): + name = 'subleft' + + def __init__(self, S_O): + super().__init__() + self.S_O = S_O + + def define(self, dimensions): + x, y = dimensions + return {x: ('left', self.S_O//2), y: ('middle', 1, 1)} + + +class SubRight(SubDomain): + name = 'subright' + + def __init__(self, S_O): + super().__init__() + self.S_O = S_O + + def define(self, dimensions): + x, y = dimensions + return {x: ('right', self.S_O//2), y: ('middle', 1, 1)} + + +class SubPointBottomLeft(SubDomain): + name = 'subpointbottomleft' + + def define(self, dimensions): + x, y = dimensions + return {x: ('left', 1), y: ('left', 1)} + + +class SubPointBottomRight(SubDomain): + name = 'subpointbottomright' + + def define(self, dimensions): + x, y = dimensions + return {x: ('right', 1), y: ('left', 1)} + + +class SubPointTopLeft(SubDomain): + name = 'subpointtopleft' + + def define(self, dimensions): + x, y = dimensions + return {x: ('left', 1), y: ('right', 1)} + + +class SubPointTopRight(SubDomain): + name = 'subpointtopright' + + def define(self, dimensions): + x, y = dimensions + return {x: ('right', 1), y: ('right', 1)} + + +def neumann_bottom(eq, subdomain): + lhs, rhs = eq.evaluate.args + + # Get vertical subdimension and its parent + yfs = subdomain.dimensions[-1] + y = yfs.parent + + # Functions present in stencil + funcs = retrieve_functions(lhs-rhs) + + mapper = {} + for f in funcs: + # Get the y index + yind = f.indices[-1] + if (yind - y).as_coeff_Mul()[0] < 0: + if f.name == 'pn1': + mapper.update({f: f.subs({yind: INT(abs(yind))})}) + + return Eq(lhs.subs(mapper), rhs.subs(mapper), subdomain=subdomain) + + +def neumann_top(eq, subdomain): + lhs, rhs = eq.evaluate.args + + # Get vertical subdimension and its parent + yfs = subdomain.dimensions[-1] + y = yfs.parent + + # Functions present in stencil + funcs = retrieve_functions(lhs-rhs) + + mapper = {} + for f in funcs: + # Get the y index + yind = f.indices[-1] + if (yind - y).as_coeff_Mul()[0] > 0: + # Symmetric mirror + tmp = y - INT(abs(y.symbolic_max - yind)) + if f.name == 'pn1': + mapper.update({f: f.subs({yind: tmp})}) + + return Eq(lhs.subs(mapper), rhs.subs(mapper), subdomain=subdomain) + + +def neumann_left(eq, subdomain): + lhs, rhs = eq.evaluate.args + + # Get horizontal subdimension and its parent + xfs = subdomain.dimensions[0] + x = xfs.parent + + # Functions present in stencil + funcs = retrieve_functions(lhs-rhs) + + mapper = {} + for f in funcs: + # Get the x index + xind = f.indices[-2] + if (xind - x).as_coeff_Mul()[0] < 0: + # Symmetric mirror + # Substitute where index is negative for +ve + # where index is positive + if f.name == 'pn1': + mapper.update({f: f.subs({xind: INT(abs(xind))})}) + + return Eq(lhs.subs(mapper), rhs.subs(mapper), subdomain=subdomain) + + +def neumann_right(eq, subdomain): + lhs, rhs = eq.evaluate.args + + # Get horizontal subdimension and its parent + xfs = subdomain.dimensions[0] + x = xfs.parent + + # Functions present in stencil + funcs = retrieve_functions(lhs-rhs) + + mapper = {} + for f in funcs: + # Get the x index + xind = f.indices[-2] + if (xind - x).as_coeff_Mul()[0] > 0: + tmp = x - INT(abs(x.symbolic_max - xind)) + if f.name == 'pn1': + mapper.update({f: f.subs({xind: tmp})}) + + return Eq(lhs.subs(mapper), rhs.subs(mapper), subdomain=subdomain) + + +sub1 = SubTop(so) +sub2 = SubBottom(so) +sub3 = SubLeft(so) +sub4 = SubRight(so) +sub5 = SubPointBottomLeft() +sub6 = SubPointBottomRight() +sub7 = SubPointTopLeft() +sub8 = SubPointTopRight() + +subdomains = (sub1, sub2, sub3, sub4, sub5, sub6, sub7, sub8) + +grid = Grid( + shape=(nx, ny), extent=(Lx, Ly), subdomains=subdomains, dtype=np.float64 +) +time = grid.time_dim +t = grid.stepping_dim +x, y = grid.dimensions + +# time stepping parameters +dt = 1e-3 +t_end = 1. +ns = int(t_end/dt) + +u1 = TimeFunction(name='u1', grid=grid, space_order=2, dtype=np.float64) +v1 = TimeFunction(name='v1', grid=grid, space_order=2, dtype=np.float64) +pn1 = Function(name='pn1', grid=grid, space_order=2, dtype=np.float64) + +pn1.data[:] = 0. + +eq_pn1 = Eq(pn1.laplace, rho*(1./dt*(u1.forward.dxc+v1.forward.dyc)), + subdomain=grid.interior) + + +bc_pn1 = [neumann_top(eq_pn1, sub1)] +bc_pn1 += [neumann_bottom(eq_pn1, sub2)] +bc_pn1 += [neumann_left(eq_pn1, sub3)] +bc_pn1 += [neumann_right(eq_pn1, sub4)] +bc_pn1 += [EssentialBC(pn1, 0., subdomain=sub5)] +bc_pn1 += [neumann_right(neumann_bottom(eq_pn1, sub6), sub6)] +bc_pn1 += [neumann_left(neumann_top(eq_pn1, sub7), sub7)] +bc_pn1 += [neumann_right(neumann_top(eq_pn1, sub8), sub8)] + + +eqn_p = PETScSolve([eq_pn1]+bc_pn1, pn1) + +eq_u1 = Eq(u1.dt + u1*u1.dxc + v1*u1.dyc, nu*u1.laplace) +eq_v1 = Eq(v1.dt + u1*v1.dxc + v1*v1.dyc, nu*v1.laplace) + +update_u = Eq(u1.forward, u1.forward - (dt/rho)*(pn1.dxc), + subdomain=grid.interior) + +update_v = Eq(v1.forward, v1.forward - (dt/rho)*(pn1.dyc), + subdomain=grid.interior) + +# TODO: Can drop due to initial guess CB +u1.data[0, :, -1] = np.float64(1.) +u1.data[1, :, -1] = np.float64(1.) + + +# TODO: Don't need both sets of bcs, can reuse petsc ones +# Create Dirichlet BC expressions for velocity +bc_u1 = [Eq(u1[t+1, x, ny-1], 1.)] # top +bc_u1 += [Eq(u1[t+1, 0, y], 0.)] # left +bc_u1 += [Eq(u1[t+1, nx-1, y], 0.)] # right +bc_u1 += [Eq(u1[t+1, x, 0], 0.)] # bottom +bc_v1 = [Eq(v1[t+1, 0, y], 0.)] # left +bc_v1 += [Eq(v1[t+1, nx-1, y], 0.)] # right +bc_v1 += [Eq(v1[t+1, x, ny-1], 0.)] # top +bc_v1 += [Eq(v1[t+1, x, 0], 0.)] # bottom + +# Create Dirichlet BC expressions for velocity +bc_petsc_u1 = [EssentialBC(u1.forward, 1., subdomain=sub1)] # top +bc_petsc_u1 += [EssentialBC(u1.forward, 0., subdomain=sub3)] # left +bc_petsc_u1 += [EssentialBC(u1.forward, 0., subdomain=sub4)] # right +bc_petsc_u1 += [EssentialBC(u1.forward, 0., subdomain=sub2)] # bottom +bc_petsc_v1 = [EssentialBC(v1.forward, 0., subdomain=sub3)] # left +bc_petsc_v1 += [EssentialBC(v1.forward, 0., subdomain=sub4)] # right +bc_petsc_v1 += [EssentialBC(v1.forward, 0., subdomain=sub1)] # top +bc_petsc_v1 += [EssentialBC(v1.forward, 0., subdomain=sub2)] # bottom + +tentu = PETScSolve([eq_u1]+bc_petsc_u1, u1.forward) +tentv = PETScSolve([eq_v1]+bc_petsc_v1, v1.forward) + +exprs = tentu + tentv + eqn_p + [update_u, update_v] + bc_u1 + bc_v1 + +with switchconfig(language='petsc'): + op = Operator(exprs) + op.apply(time_m=0, time_M=ns-1, dt=dt) + +u1_norm = norm(u1) +v1_norm = norm(v1) +p1_norm = norm(pn1) + + +# TODO: change these norm checks to array checks (use paper) +assert np.isclose(u1_norm, 13.966067703420883, atol=0, rtol=1e-7) +assert np.isclose(v1_norm, 7.9575677674738285, atol=0, rtol=1e-7) +assert np.isclose(p1_norm, 36.46263134701362, atol=0, rtol=1e-7) diff --git a/examples/petsc/random/01_helmholtz.py b/examples/petsc/random/01_helmholtz.py new file mode 100644 index 0000000000..8702fbf298 --- /dev/null +++ b/examples/petsc/random/01_helmholtz.py @@ -0,0 +1,259 @@ +import os +import numpy as np + +from devito.symbolics import retrieve_functions, INT +from devito import (configuration, Operator, Eq, Grid, Function, + SubDomain, switchconfig) +from devito.petsc import PETScSolve +from devito.petsc.initialize import PetscInitialize +configuration['compiler'] = 'custom' +os.environ['CC'] = 'mpicc' + +# Modified Helmholtz equation +# Ref - https://www.firedrakeproject.org/demos/helmholtz.py.html + +PetscInitialize() + + +Lx = 1. +Ly = Lx + +# Number of grid points in each direction +n = 11 + +# mesh spacing +dx = Lx/(n-1) +dy = Ly/(n-1) +so = 2 + + +class SubTop(SubDomain): + name = 'subtop' + + def __init__(self, S_O): + super().__init__() + self.S_O = S_O + + def define(self, dimensions): + x, y = dimensions + return {x: ('middle', 1, 1), y: ('right', self.S_O//2)} + + +class SubBottom(SubDomain): + name = 'subbottom' + + def __init__(self, S_O): + super().__init__() + self.S_O = S_O + + def define(self, dimensions): + x, y = dimensions + return {x: ('middle', 1, 1), y: ('left', self.S_O//2)} + + +class SubLeft(SubDomain): + name = 'subleft' + + def __init__(self, S_O): + super().__init__() + self.S_O = S_O + + def define(self, dimensions): + x, y = dimensions + return {x: ('left', self.S_O//2), y: ('middle', 1, 1)} + + +class SubRight(SubDomain): + name = 'subright' + + def __init__(self, S_O): + super().__init__() + self.S_O = S_O + + def define(self, dimensions): + x, y = dimensions + return {x: ('right', self.S_O//2), y: ('middle', 1, 1)} + + +class SubPointBottomLeft(SubDomain): + name = 'subpointbottomleft' + + def define(self, dimensions): + x, y = dimensions + return {x: ('left', 1), y: ('left', 1)} + + +class SubPointBottomRight(SubDomain): + name = 'subpointbottomright' + + def define(self, dimensions): + x, y = dimensions + return {x: ('right', 1), y: ('left', 1)} + + +class SubPointTopLeft(SubDomain): + name = 'subpointtopleft' + + def define(self, dimensions): + x, y = dimensions + return {x: ('left', 1), y: ('right', 1)} + + +class SubPointTopRight(SubDomain): + name = 'subpointtopright' + + def define(self, dimensions): + x, y = dimensions + return {x: ('right', 1), y: ('right', 1)} + + +def neumann_bottom(eq, subdomain): + lhs, rhs = eq.evaluate.args + + # Get vertical subdimension and its parent + yfs = subdomain.dimensions[-1] + y = yfs.parent + + # Functions present in stencil + funcs = retrieve_functions(lhs-rhs) + + mapper = {} + for f in funcs: + # Get the y index + yind = f.indices[-1] + if (yind - y).as_coeff_Mul()[0] < 0: + mapper.update({f: f.subs({yind: INT(abs(yind))})}) + + return Eq(lhs.subs(mapper), rhs.subs(mapper), subdomain=subdomain) + + +def neumann_top(eq, subdomain): + lhs, rhs = eq.evaluate.args + + # Get vertical subdimension and its parent + yfs = subdomain.dimensions[-1] + y = yfs.parent + + # Functions present in stencil + funcs = retrieve_functions(lhs-rhs) + + mapper = {} + for f in funcs: + # Get the y index + yind = f.indices[-1] + if (yind - y).as_coeff_Mul()[0] > 0: + # Symmetric mirror + tmp = y - INT(abs(y.symbolic_max - yind)) + mapper.update({f: f.subs({yind: tmp})}) + + return Eq(lhs.subs(mapper), rhs.subs(mapper), subdomain=subdomain) + + +def neumann_left(eq, subdomain): + lhs, rhs = eq.evaluate.args + + # Get horizontal subdimension and its parent + xfs = subdomain.dimensions[0] + x = xfs.parent + + # Functions present in stencil + funcs = retrieve_functions(lhs-rhs) + + mapper = {} + for f in funcs: + # Get the x index + xind = f.indices[-2] + if (xind - x).as_coeff_Mul()[0] < 0: + # Symmetric mirror + # Substitute where index is negative for +ve where + # index is positive + mapper.update({f: f.subs({xind: INT(abs(xind))})}) + + return Eq(lhs.subs(mapper), rhs.subs(mapper), subdomain=subdomain) + + +def neumann_right(eq, subdomain): + lhs, rhs = eq.evaluate.args + + # Get horizontal subdimension and its parent + xfs = subdomain.dimensions[0] + x = xfs.parent + + # Functions present in stencil + funcs = retrieve_functions(lhs-rhs) + + mapper = {} + for f in funcs: + # Get the x index + xind = f.indices[-2] + if (xind - x).as_coeff_Mul()[0] > 0: + tmp = x - INT(abs(x.symbolic_max - xind)) + mapper.update({f: f.subs({xind: tmp})}) + + return Eq(lhs.subs(mapper), rhs.subs(mapper), subdomain=subdomain) + + +sub1 = SubTop(so) +sub2 = SubBottom(so) +sub3 = SubLeft(so) +sub4 = SubRight(so) +sub5 = SubPointBottomLeft() +sub6 = SubPointBottomRight() +sub7 = SubPointTopLeft() +sub8 = SubPointTopRight() + +subdomains = (sub1, sub2, sub3, sub4, sub5, sub6, sub7, sub8) + + +def analytical_solution(x, y): + return np.cos(2*np.pi*x)*np.cos(2*np.pi*y) + + +n_values = [11, 21, 31, 41, 51, 61, 71, 81, 91, 101] +h = np.array([Lx/(n-1) for n in n_values]) +errors = [] + + +for n in n_values: + grid = Grid( + shape=(n, n), extent=(Lx, Ly), subdomains=subdomains, dtype=np.float64 + ) + time = grid.time_dim + t = grid.stepping_dim + x, y = grid.dimensions + + u = Function(name='u', grid=grid, space_order=so, dtype=np.float64) + f = Function(name='f', grid=grid, space_order=so, dtype=np.float64) + + tmpx = np.linspace(0, Lx, n).astype(np.float64) + tmpy = np.linspace(0, Ly, n).astype(np.float64) + Y, X = np.meshgrid(tmpx, tmpy) + f.data[:] = (1.+(8.*(np.pi**2)))*np.cos(2.*np.pi*X)*np.cos(2.*np.pi*Y) + + eqn = Eq(-u.laplace+u, f, subdomain=grid.interior) + + bcs = [neumann_top(eqn, sub1)] + bcs += [neumann_bottom(eqn, sub2)] + bcs += [neumann_left(eqn, sub3)] + bcs += [neumann_right(eqn, sub4)] + bcs += [neumann_left(neumann_bottom(eqn, sub5), sub5)] + bcs += [neumann_right(neumann_bottom(eqn, sub6), sub6)] + bcs += [neumann_left(neumann_top(eqn, sub7), sub7)] + bcs += [neumann_right(neumann_top(eqn, sub8), sub8)] + + solver = PETScSolve([eqn]+bcs, target=u, solver_parameters={'rtol': 1e-8}) + + with switchconfig(openmp=False, language='petsc'): + op = Operator(solver) + op.apply() + + analytical = analytical_solution(X, Y) + + diff = analytical[:] - u.data[:] + error = np.linalg.norm(diff) / np.linalg.norm(analytical[:]) + errors.append(error) + +slope, _ = np.polyfit(np.log(h), np.log(errors), 1) + +assert slope > 1.9 +assert slope < 2.1 diff --git a/examples/petsc/seismic/staggered_acoustic.py b/examples/petsc/seismic/01_staggered_acoustic.py similarity index 100% rename from examples/petsc/seismic/staggered_acoustic.py rename to examples/petsc/seismic/01_staggered_acoustic.py From 4f3bd8827cd9556b199674e8472ab0e019e8c3f4 Mon Sep 17 00:00:00 2001 From: ZoeLeibowitz <90093761+ZoeLeibowitz@users.noreply.github.com> Date: Thu, 3 Apr 2025 15:00:16 +0100 Subject: [PATCH 15/84] Use VecCreateMPIWithArray (#77) --- devito/ir/iet/visitors.py | 4 +- devito/petsc/iet/routines.py | 161 +++++++++++++++++++++++++++-------- devito/petsc/types/object.py | 2 +- tests/test_symbolics.py | 2 + 4 files changed, 131 insertions(+), 38 deletions(-) diff --git a/devito/ir/iet/visitors.py b/devito/ir/iet/visitors.py index 34ac76cf06..d936b743e4 100644 --- a/devito/ir/iet/visitors.py +++ b/devito/ir/iet/visitors.py @@ -214,11 +214,11 @@ def _gen_struct_decl(self, obj, masked=()): except TypeError: # E.g., `ctype` is of type `dtypes_lowering.CustomDtype` if isinstance(obj, LocalCompositeObject): - # TODO: Potentially re-evaluate: Setting ctype to obj allows + # TODO: re-evaluate: Setting ctype to obj allows # _gen_struct_decl to generate a cgen.Structure from a # LocalCompositeObject, where obj._C_ctype is a CustomDtype. # LocalCompositeObject has a __fields__ property, - # which allows the subsequent code in this function to function + # which allows the subsequent code in this function to work # correctly. ctype = obj else: diff --git a/devito/petsc/iet/routines.py b/devito/petsc/iet/routines.py index 5a000e57a6..68765d7801 100644 --- a/devito/petsc/iet/routines.py +++ b/devito/petsc/iet/routines.py @@ -1,5 +1,6 @@ from collections import OrderedDict from functools import cached_property +import math from devito.ir.iet import (Call, FindSymbols, List, Uxreplace, CallableBody, Dereference, DummyExpr, BlankLine, Callable, FindNodes, @@ -1111,8 +1112,19 @@ def _setup(self): global_x = petsc_call('DMCreateGlobalVector', [dmda, Byref(sobjs['xglobal'])]) - local_x = petsc_call('DMCreateLocalVector', - [dmda, Byref(sobjs['xlocal'])]) + target = self.fielddata.target + field_from_ptr = FieldFromPointer( + target.function._C_field_data, target.function._C_symbol + ) + + local_size = math.prod( + v for v, dim in zip(target.shape_allocated, target.dimensions) if dim.is_Space + ) + local_x = petsc_call('VecCreateMPIWithArray', + ['PETSC_COMM_WORLD', 1, local_size, 'PETSC_DECIDE', + field_from_ptr, Byref(sobjs['xlocal'])]) + + # TODO: potentially also need to set the DM and local/global map to xlocal get_local_size = petsc_call('VecGetSize', [sobjs['xlocal'], Byref(sobjs['localsize'])]) @@ -1247,11 +1259,87 @@ class CoupledSetup(BaseSetup): def snes_ctx(self): return Byref(self.solver_objs['jacctx']) - def _extend_setup(self): + def _setup(self): + # TODO: minimise code duplication with superclass objs = self.objs sobjs = self.solver_objs dmda = sobjs['dmda'] + + solver_params = self.injectsolve.expr.rhs.solver_parameters + + snes_create = petsc_call('SNESCreate', [objs['comm'], Byref(sobjs['snes'])]) + + snes_set_dm = petsc_call('SNESSetDM', [sobjs['snes'], dmda]) + + create_matrix = petsc_call('DMCreateMatrix', [dmda, Byref(sobjs['Jac'])]) + + # NOTE: Assuming all solves are linear for now + snes_set_type = petsc_call('SNESSetType', [sobjs['snes'], 'SNESKSPONLY']) + + snes_set_jac = petsc_call( + 'SNESSetJacobian', [sobjs['snes'], sobjs['Jac'], + sobjs['Jac'], 'MatMFFDComputeJacobian', objs['Null']] + ) + + global_x = petsc_call('DMCreateGlobalVector', + [dmda, Byref(sobjs['xglobal'])]) + + local_x = petsc_call('DMCreateLocalVector', [dmda, Byref(sobjs['xlocal'])]) + + get_local_size = petsc_call('VecGetSize', + [sobjs['xlocal'], Byref(sobjs['localsize'])]) + + global_b = petsc_call('DMCreateGlobalVector', + [dmda, Byref(sobjs['bglobal'])]) + + snes_get_ksp = petsc_call('SNESGetKSP', + [sobjs['snes'], Byref(sobjs['ksp'])]) + + ksp_set_tols = petsc_call( + 'KSPSetTolerances', [sobjs['ksp'], solver_params['ksp_rtol'], + solver_params['ksp_atol'], solver_params['ksp_divtol'], + solver_params['ksp_max_it']] + ) + + ksp_set_type = petsc_call( + 'KSPSetType', [sobjs['ksp'], solver_mapper[solver_params['ksp_type']]] + ) + + ksp_get_pc = petsc_call( + 'KSPGetPC', [sobjs['ksp'], Byref(sobjs['pc'])] + ) + + # Even though the default will be jacobi, set to PCNONE for now + pc_set_type = petsc_call('PCSetType', [sobjs['pc'], 'PCNONE']) + + ksp_set_from_ops = petsc_call('KSPSetFromOptions', [sobjs['ksp']]) + + matvec = self.cbbuilder.main_matvec_callback + matvec_operation = petsc_call( + 'MatShellSetOperation', + [sobjs['Jac'], 'MATOP_MULT', MatShellSetOp(matvec.name, void, void)] + ) + formfunc = self.cbbuilder.main_formfunc_callback + formfunc_operation = petsc_call( + 'SNESSetFunction', + [sobjs['snes'], objs['Null'], FormFunctionCallback(formfunc.name, void, void), + self.snes_ctx] + ) + + dmda_calls = self._create_dmda_calls(dmda) + + mainctx = sobjs['userctx'] + + call_struct_callback = petsc_call( + self.cbbuilder.user_struct_callback.name, [Byref(mainctx)] + ) + + # TODO: maybe don't need to explictly set this + mat_set_dm = petsc_call('MatSetDM', [sobjs['Jac'], dmda]) + + calls_set_app_ctx = petsc_call('DMSetApplicationContext', [dmda, Byref(mainctx)]) + create_field_decomp = petsc_call( 'DMCreateFieldDecomposition', [dmda, Byref(sobjs['nfields']), objs['Null'], Byref(sobjs['fields']), @@ -1297,13 +1385,34 @@ def _extend_setup(self): [sobjs[f'da{t.name}'], Byref(sobjs[f'bglobal{t.name}'])] ) for t in targets] - return ( + coupled_setup = dmda_calls + ( + snes_create, + snes_set_dm, + create_matrix, + snes_set_jac, + snes_set_type, + global_x, + local_x, + get_local_size, + global_b, + snes_get_ksp, + ksp_set_tols, + ksp_set_type, + ksp_get_pc, + pc_set_type, + ksp_set_from_ops, + matvec_operation, + formfunc_operation, + call_struct_callback, + mat_set_dm, + calls_set_app_ctx, create_field_decomp, matop_create_submats_op, call_coupled_struct_callback, shell_set_ctx, - create_submats - ) + tuple(deref_dms) + tuple(xglobals) + tuple(bglobals) + create_submats) + \ + tuple(deref_dms) + tuple(xglobals) + tuple(bglobals) + return coupled_setup class Solver: @@ -1333,7 +1442,7 @@ def _execute_solve(self): rhs_call = petsc_call(rhs_callback.name, [sobjs['dmda'], sobjs['bglobal']]) - vec_replace_array = self.timedep.replace_array(target) + vec_place_array = self.timedep.place_array(target) if self.cbbuilder.initialguesses: initguess = self.cbbuilder.initialguesses[0] @@ -1358,7 +1467,7 @@ def _execute_solve(self): run_solver_calls = (struct_assignment,) + ( rhs_call, - ) + vec_replace_array + ( + ) + vec_place_array + ( initguess_call, dm_local_to_global_x, snes_solve, @@ -1415,7 +1524,7 @@ def _execute_solve(self): pre_solve += ( petsc_call(c.name, [dm, target_bglob]), petsc_call('DMCreateLocalVector', [dm, Byref(target_xloc)]), - self.timedep.replace_array(t), + self.timedep.place_array(t), petsc_call( 'DMLocalToGlobal', [dm, target_xloc, insert_vals, target_xglob] @@ -1485,23 +1594,7 @@ def _origin_to_moddim_mapper(self, iters): def uxreplace_time(self, body): return body - def replace_array(self, target): - """ - VecReplaceArray() is a PETSc function that allows replacing the array - of a `Vec` with a user provided array. - https://petsc.org/release/manualpages/Vec/VecReplaceArray/ - - This function is used to replace the array of the PETSc solution `Vec` - with the array from the `Function` object representing the target. - - Examples - -------- - >>> target - f1(x, y) - >>> call = replace_array(target) - >>> print(call) - PetscCall(VecReplaceArray(xlocal0,f1_vec->data)); - """ + def place_array(self, target): sobjs = self.sobjs field_from_ptr = FieldFromPointer( @@ -1608,29 +1701,27 @@ def _origin_to_moddim_mapper(self, iters): mapper[d] = d return mapper - def replace_array(self, target): + def place_array(self, target): """ In the case that the actual target is time-dependent e.g a `TimeFunction`, a pointer to the first element in the array that will be updated during - the time step is passed to VecReplaceArray(). + the time step is passed to VecPlaceArray(). Examples -------- >>> target f1(time + dt, x, y) - >>> calls = replace_array(target) + >>> calls = place_array(target) >>> print(List(body=calls)) - PetscCall(VecGetSize(xlocal0,&(localsize0))); float * f1_ptr0 = (time + 1)*localsize0 + (float*)(f1_vec->data); - PetscCall(VecReplaceArray(xlocal0,f1_ptr0)); + PetscCall(VecPlaceArray(xlocal0,f1_ptr0)); >>> target f1(t + dt, x, y) - >>> calls = replace_array(target) + >>> calls = place_array(target) >>> print(List(body=calls)) - PetscCall(VecGetSize(xlocal0,&(localsize0))); float * f1_ptr0 = t1*localsize0 + (float*)(f1_vec->data); - PetscCall(VecReplaceArray(xlocal0,f1_ptr0)); + PetscCall(VecPlaceArray(xlocal0,f1_ptr0)); """ sobjs = self.sobjs @@ -1654,7 +1745,7 @@ def replace_array(self, target): ), petsc_call('VecPlaceArray', [xlocal, start_ptr]) ) - return super().replace_array(target) + return super().place_array(target) def assign_time_iters(self, struct): """ diff --git a/devito/petsc/types/object.py b/devito/petsc/types/object.py index fcbf06ea69..2e69e4a0fc 100644 --- a/devito/petsc/types/object.py +++ b/devito/petsc/types/object.py @@ -36,7 +36,7 @@ def dofs(self): def _C_free(self): return petsc_call('DMDestroy', [Byref(self.function)]) - # TODO: This is growing out of hand so switch to an enumeration or something? + # TODO: Switch to an enumeration? @property def _C_free_priority(self): return 4 diff --git a/tests/test_symbolics.py b/tests/test_symbolics.py index afa32577b3..82e85f2c59 100644 --- a/tests/test_symbolics.py +++ b/tests/test_symbolics.py @@ -319,6 +319,8 @@ def test_extended_sympy_arithmetic(): o = Object(name='o', dtype=c_void_p) bar = FieldFromPointer('bar', o) # TODO: Edit/fix/update according to PR #2513 + # The order changed due to adding the dtype property + # to FieldFromPointer assert ccode(-1 + bar) == 'o->bar - 1' From 9c60990b3a79152f1f4eedce7db23b3c84c848e0 Mon Sep 17 00:00:00 2001 From: ZoeLeibowitz Date: Thu, 3 Apr 2025 21:39:58 +0100 Subject: [PATCH 16/84] workflows: Edit petsc docker image --- .github/workflows/pytest-petsc.yml | 4 ---- docker/Dockerfile.cpu | 4 ++-- docker/Dockerfile.petsc | 5 ++++- 3 files changed, 6 insertions(+), 7 deletions(-) diff --git a/.github/workflows/pytest-petsc.yml b/.github/workflows/pytest-petsc.yml index 0d44614b16..eb6e3f7965 100644 --- a/.github/workflows/pytest-petsc.yml +++ b/.github/workflows/pytest-petsc.yml @@ -56,10 +56,6 @@ jobs: run: | docker build -f docker/Dockerfile.devito --build-arg base=zoeleibowitz/petsc_image:latest --tag zoeleibowitz/petsc_devito_image:latest . - - name: Push Docker image to DockerHub - run: | - docker push ${{ secrets.DOCKER_USERNAME }}/petsc_devito_image:latest - - name: Set run prefix run: | echo "RUN_CMD=docker run --rm -t -e CODECOV_TOKEN=${{ secrets.CODECOV_TOKEN }} --name testrun zoeleibowitz/petsc_devito_image:latest" >> $GITHUB_ENV diff --git a/docker/Dockerfile.cpu b/docker/Dockerfile.cpu index 04da0dcc8d..8f036badb9 100644 --- a/docker/Dockerfile.cpu +++ b/docker/Dockerfile.cpu @@ -17,8 +17,7 @@ RUN apt-get update && \ # Install for basic base not containing it RUN apt-get install -y wget flex libnuma-dev hwloc curl cmake git \ - autoconf libtool build-essential procps software-properties-common \ - gfortran pkgconf libopenblas-serial-dev + autoconf libtool build-essential procps software-properties-common # Install gcc RUN if [ -n "$gcc" ]; then \ @@ -44,6 +43,7 @@ RUN cd /tmp && mkdir openmpi && \ cd openmpi && ./autogen.pl && \ mkdir build && cd build && \ ../configure --prefix=/opt/openmpi/ \ + --disable-mpi-fortran \ --enable-mca-no-build=btl-uct --enable-mpi1-compatibility && \ make -j ${nproc} && \ make install && \ diff --git a/docker/Dockerfile.petsc b/docker/Dockerfile.petsc index 6bf0800a8e..9d5f347d68 100644 --- a/docker/Dockerfile.petsc +++ b/docker/Dockerfile.petsc @@ -7,7 +7,10 @@ # use the one from my dockerhub ARG base=zoeleibowitz/bases:cpu-gcc -RUN apt-get update && apt-get install -y git && \ +FROM $base + +RUN apt-get update && apt-get install -y \ + git gfortran pkgconf libopenblas-serial-dev && \ python3 -m venv /venv && \ /venv/bin/pip install --no-cache-dir --upgrade pip && \ /venv/bin/pip install --no-cache-dir --no-binary numpy numpy && \ From a8f11a7aef1c6fc02d2a66c9bed86b5b00048ba7 Mon Sep 17 00:00:00 2001 From: ZoeLeibowitz Date: Thu, 3 Apr 2025 21:41:33 +0100 Subject: [PATCH 17/84] workflows: Update petsc install to latest release --- docker/Dockerfile.petsc | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/docker/Dockerfile.petsc b/docker/Dockerfile.petsc index 9d5f347d68..056c860688 100644 --- a/docker/Dockerfile.petsc +++ b/docker/Dockerfile.petsc @@ -16,7 +16,7 @@ RUN apt-get update && apt-get install -y \ /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 && \ + git clone -b v3.23.0 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) \ From d06b2752998908bb6fa8d571448f11887d2ac5b1 Mon Sep 17 00:00:00 2001 From: Jack Betteridge Date: Wed, 16 Apr 2025 17:43:16 +0100 Subject: [PATCH 18/84] logging: Corner case for PETSc --- devito/logger.py | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/devito/logger.py b/devito/logger.py index e6d1793347..9c1dfa56d2 100644 --- a/devito/logger.py +++ b/devito/logger.py @@ -134,6 +134,12 @@ def log(msg, level=INFO, *args, **kwargs): ERROR, CRITICAL``. """ color = COLORS[level] if sys.stdout.isatty() and sys.stderr.isatty() else '%s' + # TODO: Think about a proper fix for this: + # When running with PETSc, `PetscFinalize` will helpfully close the stream + # that is being used for logging before all messages are displayed. Specifically + # the operator `info` logging. + if logger.handlers[0] is stream_handler and logger.handlers[0].stream.closed: + logger.handlers[0].stream = sys.stdout logger.log(level, color % msg, *args, **kwargs) From 2b6c03883ecfc74df6872656d5ca20b2031e4c7c Mon Sep 17 00:00:00 2001 From: ZoeLeibowitz Date: Thu, 17 Apr 2025 10:39:48 +0100 Subject: [PATCH 19/84] examples: Add 1D example --- examples/petsc/Poisson/03_poisson.py | 96 ++++++++++++++++++++++++++++ 1 file changed, 96 insertions(+) create mode 100644 examples/petsc/Poisson/03_poisson.py diff --git a/examples/petsc/Poisson/03_poisson.py b/examples/petsc/Poisson/03_poisson.py new file mode 100644 index 0000000000..9b93fd31e1 --- /dev/null +++ b/examples/petsc/Poisson/03_poisson.py @@ -0,0 +1,96 @@ +import os +import numpy as np + +from devito import (Grid, Function, Eq, Operator, switchconfig, + configuration, SubDomain) + +from devito.petsc import PETScSolve, EssentialBC +from devito.petsc.initialize import PetscInitialize +configuration['compiler'] = 'custom' +os.environ['CC'] = 'mpicc' + +# 1D test +# ref - An efficient implementation of fourth-order compact finite +# difference scheme for Poisson equation with Dirichlet +# boundary conditions +# https://pdf.sciencedirectassets.com/271503/1-s2.0-S0898122116X00090/1-s2.0-S0898122116300761/main.pdf?X-Amz-Security-Token=IQoJb3JpZ2luX2VjENH%2F%2F%2F%2F%2F%2F%2F%2F%2F%2FwEaCXVzLWVhc3QtMSJGMEQCIH%2F3tmUOBoYeSj%2FmE47N2yF5b5IvdQohX7JWYGRZhNFTAiAX0P7PnDVBy2D3pTjsZRFgFEh7nVzV8oxoUFevx%2FucWSqzBQhZEAUaDDA1OTAwMzU0Njg2NSIMHBuo2q1ii2daxqQiKpAFdh12uILywsya5sCFqTsXwUdb%2F4lvp5uZIigTrb18tjM8cPru5xrZDgDVYVIlT4G6L1SE05FkjWKSQ7AO24wec3y2bNKGgpC%2FPFbEmTv6CgpR%2FjoSpToblGLiOqTkUgICSK3EoMdah%2BWl552nO0Ajdwuor0brGfDM7C2fgH1FqM%2BLyJ2do33rYFjGAswzZsQOGcf%2BChdkaKA9bxvfgE2%2Bukf7RcYDsGteSEX5Zb9XoyvMheiMUZoZk7KVPWjj3JORx9qetLs9LkpPO3IU%2BqPxtM7Vt3BnEnXR9gQ2bnL%2FtcT%2FcvsZ7a8AdiU1j8%2F%2Fxi9nBgPow0MQTmaoe9n67XRS0BVE7wAWldDb2qdZuOfwYl%2F2iG78mMTn%2FC4YcOCezc4nUT9fTcTcv3wKZzA%2Bkh8Z%2BXvdTcdADCKdVaIXLylqlhEmBlwua4cGjBG0RbpvGa%2FOBk6CbZLpn7%2FLawxsVZ1U1ksGd8HGJ%2FGMYDOauM%2FhRGNWRFsXnn%2BsrPhaJ3SoirVeV3q9JVrjGT6%2FUT3W9qIDtdPP4MJae5mp6TG5fusJjkCLxLTbeXF0%2FhbwEnAA54uj3jpTsh7rXVDB%2B8skGSdMhIITz3%2ByS%2BdMqt7iEgFOWqYXGwgXLGbOqyGGz2ikth4cs1FMT4sYrA066%2BcMkE9q3l3bsFZHQMw13UPgJQp2f69JIzgHbZ%2FoCkdDYNxUutRhZ6cMitSLrIGtcAa7p%2Fevtejnw5eTz20kLNAxjB3CMUuS1H5qhxb6cSmxneilYH1WINNPjCrDPCJ3FxlKtCJo4QzIfIKogegd%2B44T78fQzt8RP7LfA%2FzjITD9bdiCYW0f81Q3O8zzL7l7RtfnLfYXAuTFh9GtAdE8D6b4F2pnXkMwrfCCwAY6sgG4%2BnyhdUNH%2FhdcK7GZ56erHPDOYF04vpG2hZy26v7cSnA3Xb7zrqVzkLxPdyAViJnMjzV1c8itVIHgnkLuA0C%2FPJrp3RPy0ivl9dofnd%2FLtoBkoBadnTgx2f7x4SZ62bdbWk5DJ%2FavMuOajJ%2F4tl9%2F7%2FLWoyi92xH2ZCvnT4wIIakx9ODzn2dRwSYwP20omrw5oAHK8KfXr39zDhQcs6FZMnWqYVxGlKHy0XIqJY8mTLeE&X-Amz-Algorithm=AWS4-HMAC-SHA256&X-Amz-Date=20250417T092301Z&X-Amz-SignedHeaders=host&X-Amz-Expires=300&X-Amz-Credential=ASIAQ3PHCVTYXYYPN3VY%2F20250417%2Fus-east-1%2Fs3%2Faws4_request&X-Amz-Signature=bfaba1511f86b81dca784d137decbc826e87f24f30b4589d02ca67033b886446&hash=5e1a53dc45c46e59555d516468c6e00966f056dd50abcad3b239f603507d92a7&host=68042c943591013ac2b2430a89b270f6af2c76d8dfd086a07176afe7c76c2c61&pii=S0898122116300761&tid=spdf-f29c66e1-0a20-4e85-99d8-adbf3bfe5f8e&sid=68bdc92a37ea6249ef2b6425bac44510ce06gxrqb&type=client&tsoh=d3d3LnNjaWVuY2VkaXJlY3QuY29t&rh=d3d3LnNjaWVuY2VkaXJlY3QuY29t&ua=1d035b54560505005402&rr=931adc7ff9ea1b7e&cc=gb + +# Solving -u.laplace = pi^2 * sin(pix), 0 1.9 +assert slope < 2.1 From ab2450bed0815d1a68378c1b7ef45449ccad40dc Mon Sep 17 00:00:00 2001 From: ZoeLeibowitz Date: Thu, 17 Apr 2025 16:04:46 +0100 Subject: [PATCH 20/84] examples: Add 04_poisson example --- examples/petsc/Poisson/03_poisson.py | 10 ++- examples/petsc/Poisson/04_poisson.py | 123 +++++++++++++++++++++++++++ 2 files changed, 129 insertions(+), 4 deletions(-) create mode 100644 examples/petsc/Poisson/04_poisson.py diff --git a/examples/petsc/Poisson/03_poisson.py b/examples/petsc/Poisson/03_poisson.py index 9b93fd31e1..694ce6a272 100644 --- a/examples/petsc/Poisson/03_poisson.py +++ b/examples/petsc/Poisson/03_poisson.py @@ -10,10 +10,6 @@ os.environ['CC'] = 'mpicc' # 1D test -# ref - An efficient implementation of fourth-order compact finite -# difference scheme for Poisson equation with Dirichlet -# boundary conditions -# https://pdf.sciencedirectassets.com/271503/1-s2.0-S0898122116X00090/1-s2.0-S0898122116300761/main.pdf?X-Amz-Security-Token=IQoJb3JpZ2luX2VjENH%2F%2F%2F%2F%2F%2F%2F%2F%2F%2FwEaCXVzLWVhc3QtMSJGMEQCIH%2F3tmUOBoYeSj%2FmE47N2yF5b5IvdQohX7JWYGRZhNFTAiAX0P7PnDVBy2D3pTjsZRFgFEh7nVzV8oxoUFevx%2FucWSqzBQhZEAUaDDA1OTAwMzU0Njg2NSIMHBuo2q1ii2daxqQiKpAFdh12uILywsya5sCFqTsXwUdb%2F4lvp5uZIigTrb18tjM8cPru5xrZDgDVYVIlT4G6L1SE05FkjWKSQ7AO24wec3y2bNKGgpC%2FPFbEmTv6CgpR%2FjoSpToblGLiOqTkUgICSK3EoMdah%2BWl552nO0Ajdwuor0brGfDM7C2fgH1FqM%2BLyJ2do33rYFjGAswzZsQOGcf%2BChdkaKA9bxvfgE2%2Bukf7RcYDsGteSEX5Zb9XoyvMheiMUZoZk7KVPWjj3JORx9qetLs9LkpPO3IU%2BqPxtM7Vt3BnEnXR9gQ2bnL%2FtcT%2FcvsZ7a8AdiU1j8%2F%2Fxi9nBgPow0MQTmaoe9n67XRS0BVE7wAWldDb2qdZuOfwYl%2F2iG78mMTn%2FC4YcOCezc4nUT9fTcTcv3wKZzA%2Bkh8Z%2BXvdTcdADCKdVaIXLylqlhEmBlwua4cGjBG0RbpvGa%2FOBk6CbZLpn7%2FLawxsVZ1U1ksGd8HGJ%2FGMYDOauM%2FhRGNWRFsXnn%2BsrPhaJ3SoirVeV3q9JVrjGT6%2FUT3W9qIDtdPP4MJae5mp6TG5fusJjkCLxLTbeXF0%2FhbwEnAA54uj3jpTsh7rXVDB%2B8skGSdMhIITz3%2ByS%2BdMqt7iEgFOWqYXGwgXLGbOqyGGz2ikth4cs1FMT4sYrA066%2BcMkE9q3l3bsFZHQMw13UPgJQp2f69JIzgHbZ%2FoCkdDYNxUutRhZ6cMitSLrIGtcAa7p%2Fevtejnw5eTz20kLNAxjB3CMUuS1H5qhxb6cSmxneilYH1WINNPjCrDPCJ3FxlKtCJo4QzIfIKogegd%2B44T78fQzt8RP7LfA%2FzjITD9bdiCYW0f81Q3O8zzL7l7RtfnLfYXAuTFh9GtAdE8D6b4F2pnXkMwrfCCwAY6sgG4%2BnyhdUNH%2FhdcK7GZ56erHPDOYF04vpG2hZy26v7cSnA3Xb7zrqVzkLxPdyAViJnMjzV1c8itVIHgnkLuA0C%2FPJrp3RPy0ivl9dofnd%2FLtoBkoBadnTgx2f7x4SZ62bdbWk5DJ%2FavMuOajJ%2F4tl9%2F7%2FLWoyi92xH2ZCvnT4wIIakx9ODzn2dRwSYwP20omrw5oAHK8KfXr39zDhQcs6FZMnWqYVxGlKHy0XIqJY8mTLeE&X-Amz-Algorithm=AWS4-HMAC-SHA256&X-Amz-Date=20250417T092301Z&X-Amz-SignedHeaders=host&X-Amz-Expires=300&X-Amz-Credential=ASIAQ3PHCVTYXYYPN3VY%2F20250417%2Fus-east-1%2Fs3%2Faws4_request&X-Amz-Signature=bfaba1511f86b81dca784d137decbc826e87f24f30b4589d02ca67033b886446&hash=5e1a53dc45c46e59555d516468c6e00966f056dd50abcad3b239f603507d92a7&host=68042c943591013ac2b2430a89b270f6af2c76d8dfd086a07176afe7c76c2c61&pii=S0898122116300761&tid=spdf-f29c66e1-0a20-4e85-99d8-adbf3bfe5f8e&sid=68bdc92a37ea6249ef2b6425bac44510ce06gxrqb&type=client&tsoh=d3d3LnNjaWVuY2VkaXJlY3QuY29t&rh=d3d3LnNjaWVuY2VkaXJlY3QuY29t&ua=1d035b54560505005402&rr=931adc7ff9ea1b7e&cc=gb # Solving -u.laplace = pi^2 * sin(pix), 0 1.9 assert slope < 2.1 + + +# ref - An efficient implementation of fourth-order compact finite +# difference scheme for Poisson equation with Dirichlet +# boundary conditions +# https://pdf.sciencedirectassets.com/271503/1-s2.0-S0898122116X00090/1-s2.0-S0898122116300761/main.pdf?X-Amz-Security-Token=IQoJb3JpZ2luX2VjENH%2F%2F%2F%2F%2F%2F%2F%2F%2F%2FwEaCXVzLWVhc3QtMSJGMEQCIH%2F3tmUOBoYeSj%2FmE47N2yF5b5IvdQohX7JWYGRZhNFTAiAX0P7PnDVBy2D3pTjsZRFgFEh7nVzV8oxoUFevx%2FucWSqzBQhZEAUaDDA1OTAwMzU0Njg2NSIMHBuo2q1ii2daxqQiKpAFdh12uILywsya5sCFqTsXwUdb%2F4lvp5uZIigTrb18tjM8cPru5xrZDgDVYVIlT4G6L1SE05FkjWKSQ7AO24wec3y2bNKGgpC%2FPFbEmTv6CgpR%2FjoSpToblGLiOqTkUgICSK3EoMdah%2BWl552nO0Ajdwuor0brGfDM7C2fgH1FqM%2BLyJ2do33rYFjGAswzZsQOGcf%2BChdkaKA9bxvfgE2%2Bukf7RcYDsGteSEX5Zb9XoyvMheiMUZoZk7KVPWjj3JORx9qetLs9LkpPO3IU%2BqPxtM7Vt3BnEnXR9gQ2bnL%2FtcT%2FcvsZ7a8AdiU1j8%2F%2Fxi9nBgPow0MQTmaoe9n67XRS0BVE7wAWldDb2qdZuOfwYl%2F2iG78mMTn%2FC4YcOCezc4nUT9fTcTcv3wKZzA%2Bkh8Z%2BXvdTcdADCKdVaIXLylqlhEmBlwua4cGjBG0RbpvGa%2FOBk6CbZLpn7%2FLawxsVZ1U1ksGd8HGJ%2FGMYDOauM%2FhRGNWRFsXnn%2BsrPhaJ3SoirVeV3q9JVrjGT6%2FUT3W9qIDtdPP4MJae5mp6TG5fusJjkCLxLTbeXF0%2FhbwEnAA54uj3jpTsh7rXVDB%2B8skGSdMhIITz3%2ByS%2BdMqt7iEgFOWqYXGwgXLGbOqyGGz2ikth4cs1FMT4sYrA066%2BcMkE9q3l3bsFZHQMw13UPgJQp2f69JIzgHbZ%2FoCkdDYNxUutRhZ6cMitSLrIGtcAa7p%2Fevtejnw5eTz20kLNAxjB3CMUuS1H5qhxb6cSmxneilYH1WINNPjCrDPCJ3FxlKtCJo4QzIfIKogegd%2B44T78fQzt8RP7LfA%2FzjITD9bdiCYW0f81Q3O8zzL7l7RtfnLfYXAuTFh9GtAdE8D6b4F2pnXkMwrfCCwAY6sgG4%2BnyhdUNH%2FhdcK7GZ56erHPDOYF04vpG2hZy26v7cSnA3Xb7zrqVzkLxPdyAViJnMjzV1c8itVIHgnkLuA0C%2FPJrp3RPy0ivl9dofnd%2FLtoBkoBadnTgx2f7x4SZ62bdbWk5DJ%2FavMuOajJ%2F4tl9%2F7%2FLWoyi92xH2ZCvnT4wIIakx9ODzn2dRwSYwP20omrw5oAHK8KfXr39zDhQcs6FZMnWqYVxGlKHy0XIqJY8mTLeE&X-Amz-Algorithm=AWS4-HMAC-SHA256&X-Amz-Date=20250417T092301Z&X-Amz-SignedHeaders=host&X-Amz-Expires=300&X-Amz-Credential=ASIAQ3PHCVTYXYYPN3VY%2F20250417%2Fus-east-1%2Fs3%2Faws4_request&X-Amz-Signature=bfaba1511f86b81dca784d137decbc826e87f24f30b4589d02ca67033b886446&hash=5e1a53dc45c46e59555d516468c6e00966f056dd50abcad3b239f603507d92a7&host=68042c943591013ac2b2430a89b270f6af2c76d8dfd086a07176afe7c76c2c61&pii=S0898122116300761&tid=spdf-f29c66e1-0a20-4e85-99d8-adbf3bfe5f8e&sid=68bdc92a37ea6249ef2b6425bac44510ce06gxrqb&type=client&tsoh=d3d3LnNjaWVuY2VkaXJlY3QuY29t&rh=d3d3LnNjaWVuY2VkaXJlY3QuY29t&ua=1d035b54560505005402&rr=931adc7ff9ea1b7e&cc=gb diff --git a/examples/petsc/Poisson/04_poisson.py b/examples/petsc/Poisson/04_poisson.py new file mode 100644 index 0000000000..a2b8a409a5 --- /dev/null +++ b/examples/petsc/Poisson/04_poisson.py @@ -0,0 +1,123 @@ +import os +import numpy as np + +from devito import (Grid, Function, Eq, Operator, switchconfig, + configuration, SubDomain) + +from devito.petsc import PETScSolve, EssentialBC +from devito.petsc.initialize import PetscInitialize +configuration['compiler'] = 'custom' +os.environ['CC'] = 'mpicc' + + +# 2D test +# Solving u.laplace = 0 +# Dirichlet BCs. +# ref - https://www.scirp.org/journal/paperinformation?paperid=113731#f2 +# example 2 -> note they wrote u(x,1) bc wrong, it should be u(x,y) = e^-pi*sin(pix) + + +PetscInitialize() + + +# Subdomains to implement BCs +class SubTop(SubDomain): + name = 'subtop' + + def define(self, dimensions): + x, y = dimensions + return {x: x, y: ('right', 1)} + + +class SubBottom(SubDomain): + name = 'subbottom' + + def define(self, dimensions): + x, y = dimensions + return {x: x, y: ('left', 1)} + + +class SubLeft(SubDomain): + name = 'subleft' + + def define(self, dimensions): + x, y = dimensions + return {x: ('left', 1), y: y} + + +class SubRight(SubDomain): + name = 'subright' + + def define(self, dimensions): + x, y = dimensions + return {x: ('right', 1), y: y} + + +sub1 = SubTop() +sub2 = SubBottom() +sub3 = SubLeft() +sub4 = SubRight() + +subdomains = (sub1, sub2, sub3, sub4) + + +def analytical(x, y): + return np.float64(np.exp(-y*np.pi)) * np.float64(np.sin(np.pi*x)) + + +Lx = np.float64(1.) +Ly = np.float64(1.) + +n_values = [10, 30, 50, 70, 90, 110] +dx = np.array([Lx/(n-1) for n in n_values]) +errors = [] + + +for n in n_values: + grid = Grid( + shape=(n, n), extent=(Lx, Ly), subdomains=subdomains, dtype=np.float64 + ) + + u = Function(name='u', grid=grid, space_order=2) + rhs = Function(name='rhs', grid=grid, space_order=2) + + eqn = Eq(rhs, u.laplace, subdomain=grid.interior) + + tmpx = np.linspace(0, Lx, n).astype(np.float64) + tmpy = np.linspace(0, Ly, n).astype(np.float64) + + Y, X = np.meshgrid(tmpx, tmpy) + + rhs.data[:] = 0. + + bcs = Function(name='bcs', grid=grid, space_order=2) + + bcs.data[:, 0] = np.sin(np.pi*tmpx) + bcs.data[:, -1] = np.exp(-np.pi)*np.sin(np.pi*tmpx) + bcs.data[0, :] = 0. + bcs.data[-1, :] = 0. + + # # Create boundary condition expressions using subdomains + bc_eqns = [EssentialBC(u, bcs, subdomain=sub1)] + bc_eqns += [EssentialBC(u, bcs, subdomain=sub2)] + bc_eqns += [EssentialBC(u, bcs, subdomain=sub3)] + bc_eqns += [EssentialBC(u, bcs, subdomain=sub4)] + + exprs = [eqn]+bc_eqns + petsc = PETScSolve(exprs, target=u, solver_parameters={'ksp_rtol': 1e-6}) + + with switchconfig(language='petsc'): + op = Operator(petsc) + op.apply() + + u_exact = analytical(X, Y) + + diff = u_exact[1:-1] - u.data[1:-1] + error = np.linalg.norm(diff) / np.linalg.norm(u_exact[1:-1]) + errors.append(error) + +slope, _ = np.polyfit(np.log(dx), np.log(errors), 1) + + +assert slope > 1.9 +assert slope < 2.1 From dd73033eb4c612d0a9d495ce8f8dda9d1d56a1cc Mon Sep 17 00:00:00 2001 From: ZoeLeibowitz Date: Thu, 17 Apr 2025 16:05:46 +0100 Subject: [PATCH 21/84] workflows: Update workflow to run new tests --- .github/workflows/pytest-petsc.yml | 2 ++ 1 file changed, 2 insertions(+) diff --git a/.github/workflows/pytest-petsc.yml b/.github/workflows/pytest-petsc.yml index eb6e3f7965..13177cd7b8 100644 --- a/.github/workflows/pytest-petsc.yml +++ b/.github/workflows/pytest-petsc.yml @@ -80,6 +80,8 @@ jobs: ${{ env.RUN_CMD }} mpiexec -n 1 python3 examples/petsc/cfd/01_navierstokes.py ${{ env.RUN_CMD }} mpiexec -n 1 python3 examples/petsc/Poisson/01_poisson.py ${{ env.RUN_CMD }} mpiexec -n 1 python3 examples/petsc/Poisson/02_laplace.py + ${{ env.RUN_CMD }} mpiexec -n 1 python3 examples/petsc/Poisson/03_poisson.py + ${{ env.RUN_CMD }} mpiexec -n 1 python3 examples/petsc/Poisson/04_poisson.py ${{ env.RUN_CMD }} mpiexec -n 1 python3 examples/petsc/random/01_helmholtz.py - name: Upload coverage to Codecov From d3e7be6e114937d4c737d6b5ddea056ac49f4ae3 Mon Sep 17 00:00:00 2001 From: ZoeLeibowitz Date: Thu, 17 Apr 2025 18:17:14 +0100 Subject: [PATCH 22/84] misc: Flake 8 --- examples/petsc/Poisson/03_poisson.py | 1 + examples/petsc/Poisson/04_poisson.py | 2 +- 2 files changed, 2 insertions(+), 1 deletion(-) diff --git a/examples/petsc/Poisson/03_poisson.py b/examples/petsc/Poisson/03_poisson.py index 694ce6a272..9fa9a9e68a 100644 --- a/examples/petsc/Poisson/03_poisson.py +++ b/examples/petsc/Poisson/03_poisson.py @@ -37,6 +37,7 @@ def define(self, dimensions): x, = dimensions return {x: ('right', 1)} + sub1 = SubLeft() sub2 = SubRight() diff --git a/examples/petsc/Poisson/04_poisson.py b/examples/petsc/Poisson/04_poisson.py index a2b8a409a5..44f34ec8f9 100644 --- a/examples/petsc/Poisson/04_poisson.py +++ b/examples/petsc/Poisson/04_poisson.py @@ -96,7 +96,7 @@ def analytical(x, y): bcs.data[:, -1] = np.exp(-np.pi)*np.sin(np.pi*tmpx) bcs.data[0, :] = 0. bcs.data[-1, :] = 0. - + # # Create boundary condition expressions using subdomains bc_eqns = [EssentialBC(u, bcs, subdomain=sub1)] bc_eqns += [EssentialBC(u, bcs, subdomain=sub2)] From 7e21e5b309b3c74b818dd135248c2f2132212894 Mon Sep 17 00:00:00 2001 From: ZoeLeibowitz Date: Mon, 21 Apr 2025 16:56:19 +0100 Subject: [PATCH 23/84] misc: Flake8 --- tests/test_symbolics.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/test_symbolics.py b/tests/test_symbolics.py index 385f1ef4cb..deedbec950 100644 --- a/tests/test_symbolics.py +++ b/tests/test_symbolics.py @@ -14,7 +14,7 @@ CallFromPointer, Cast, DefFunction, FieldFromPointer, INT, FieldFromComposite, IntDiv, Namespace, Rvalue, ReservedWord, ListInitializer, uxreplace, pow_to_mul, - retrieve_derivatives, BaseCast, SizeOf) + retrieve_derivatives, BaseCast, SizeOf, sympy_dtype) from devito.tools import as_tuple from devito.types import (Array, Bundle, FIndexed, LocalObject, Object, ComponentAccess, StencilDimension, Symbol as dSymbol, From 44b4aebda0c0e971a91ab5d8b58f47cdef398be3 Mon Sep 17 00:00:00 2001 From: ZoeLeibowitz Date: Tue, 10 Jun 2025 16:53:40 +0100 Subject: [PATCH 24/84] misc: Fix merge leftover --- devito/types/array.py | 1 + 1 file changed, 1 insertion(+) diff --git a/devito/types/array.py b/devito/types/array.py index c17fd9317a..02ecb36453 100644 --- a/devito/types/array.py +++ b/devito/types/array.py @@ -75,6 +75,7 @@ def _C_free(self): def _C_free_priority(self): return 0 + @property def c0(self): # ArrayBasic can be used as a base class for tensorial objects (that is, # arrays whose components are AbstractFunctions). This property enables From 92edf688a666ecaffe28fff4bbd8c452e9589e69 Mon Sep 17 00:00:00 2001 From: ZoeLeibowitz Date: Thu, 24 Apr 2025 18:02:30 +0100 Subject: [PATCH 25/84] dsl/compiler: Edits to mixed solver functionality --- devito/passes/iet/definitions.py | 2 +- devito/passes/iet/engine.py | 14 +- devito/passes/iet/misc.py | 3 +- devito/petsc/iet/routines.py | 226 ++++++++++++++++++++++++------- devito/petsc/solve.py | 139 +++++++++++++++---- devito/petsc/types/array.py | 80 ++++++++++- devito/petsc/types/object.py | 7 + devito/petsc/types/types.py | 60 +++++--- devito/types/basic.py | 1 + 9 files changed, 424 insertions(+), 108 deletions(-) diff --git a/devito/passes/iet/definitions.py b/devito/passes/iet/definitions.py index e56c86645f..f5ae337b9c 100644 --- a/devito/passes/iet/definitions.py +++ b/devito/passes/iet/definitions.py @@ -379,7 +379,7 @@ def place_definitions(self, iet, globs=None, **kwargs): # created by the compiler up to this point (Array, LocalObject, etc.) storage = Storage() defines = FindSymbols('defines-aliases|globals').visit(iet) - + # from IPython import embed; embed() for i in FindSymbols().visit(iet): if i in defines: continue diff --git a/devito/passes/iet/engine.py b/devito/passes/iet/engine.py index a88cf58d45..992096fd41 100644 --- a/devito/passes/iet/engine.py +++ b/devito/passes/iet/engine.py @@ -148,7 +148,10 @@ def apply(self, func, **kwargs): # Minimize code size if len(efuncs) > len(self.efuncs): efuncs = reuse_compounds(efuncs, self.sregistry) - efuncs = reuse_efuncs(self.root, efuncs, self.sregistry) + # TODO: fix for petsc bundles + # TODO: somethinng to do with this bug is causing the compiler not to combine loops with petsc bundles + # from IPython import embed; embed() + # efuncs = reuse_efuncs(self.root, efuncs, self.sregistry) self.efuncs = efuncs @@ -391,9 +394,9 @@ def abstract_efunc(efunc): - Objects are renamed as "o0", "o1", ... """ functions = FindSymbols('basics|symbolics|dimensions').visit(efunc) - + # from IPython import embed; embed() mapper = abstract_objects(functions) - + # from IPython import embed; embed() efunc = Uxreplace(mapper).visit(efunc) efunc = efunc._rebuild(name='foo') @@ -408,9 +411,10 @@ def abstract_objects(objects0, sregistry=None): objects = [] for i in objects0: if i.is_Bundle: + # from IPython import embed; embed() objects.extend(i.components) objects.append(i) - + # from IPython import embed; embed() # Precedence rules make it possible to reconstruct objects that depend on # higher priority objects keys = [Bundle, Array, DiscreteFunction, AbstractIncrDimension, BlockDimension] @@ -420,7 +424,9 @@ def abstract_objects(objects0, sregistry=None): # Build abstraction mappings mapper = {} sregistry = sregistry or SymbolRegistry() + # from IPython import embed; embed() for i in objects: + # from IPython import embed; embed() abstract_object(i, mapper, sregistry) return mapper diff --git a/devito/passes/iet/misc.py b/devito/passes/iet/misc.py index e404a8e373..0f15deec81 100644 --- a/devito/passes/iet/misc.py +++ b/devito/passes/iet/misc.py @@ -184,6 +184,7 @@ def _generate_macros_findexeds(iet, sregistry=None, tracker=None, **kwargs): try: v = tracker[i.base].v subs[i] = v.func(v.base, *i.indices) + # from IPython import embed; embed() continue except KeyError: pass @@ -193,7 +194,7 @@ def _generate_macros_findexeds(iet, sregistry=None, tracker=None, **kwargs): subs[i] = v tracker[i.base] = Bunch(header=header, v=v) - + # from IPython import embed; embed() iet = Uxreplace(subs).visit(iet) return iet diff --git a/devito/petsc/iet/routines.py b/devito/petsc/iet/routines.py index 68765d7801..be76ceace1 100644 --- a/devito/petsc/iet/routines.py +++ b/devito/petsc/iet/routines.py @@ -4,15 +4,16 @@ from devito.ir.iet import (Call, FindSymbols, List, Uxreplace, CallableBody, Dereference, DummyExpr, BlankLine, Callable, FindNodes, - retrieve_iteration_tree, filter_iterations, Iteration) + retrieve_iteration_tree, filter_iterations, Iteration, + PointerCast, Expression, Transformer) from devito.symbolics import (Byref, FieldFromPointer, cast, VOID, - FieldFromComposite, IntDiv, Deref, Mod) + FieldFromComposite, IntDiv, Deref, Mod, uxreplace) from devito.symbolics.unevaluation import Mul from devito.types.basic import AbstractFunction from devito.types import Temp, Dimension from devito.tools import filter_ordered -from devito.petsc.types import PETScArray +from devito.petsc.types import PETScArray, PetscBundle from devito.petsc.iet.nodes import (PETScCallable, FormFunctionCallback, MatShellSetOp, PetscMetaData) from devito.petsc.iet.utils import petsc_call, petsc_struct @@ -20,7 +21,7 @@ from devito.petsc.types import (DM, Mat, CallbackVec, Vec, KSP, PC, SNES, PetscInt, StartPtr, PointerIS, PointerDM, VecScatter, DMCast, JacobianStructCast, JacobianStruct, - SubMatrixStruct, CallbackDM) + SubMatrixStruct, CallbackDM, PetscComponentAccess) class CBBuilder: @@ -100,20 +101,20 @@ def user_struct_callback(self): def _make_core(self): fielddata = self.injectsolve.expr.rhs.fielddata - self._make_matvec(fielddata, fielddata.matvecs) + self._make_matvec(fielddata.arrays, fielddata.matvecs) self._make_formfunc(fielddata) self._make_formrhs(fielddata) if fielddata.initialguess: self._make_initialguess(fielddata) self._make_user_struct_callback() - def _make_matvec(self, fielddata, matvecs, prefix='MatMult'): + def _make_matvec(self, arrays, matvecs, prefix='MatMult'): # Compile matvec `eqns` into an IET via recursive compilation irs_matvec, _ = self.rcompile(matvecs, options={'mpi': False}, sregistry=self.sregistry, concretize_mapper=self.concretize_mapper) body_matvec = self._create_matvec_body(List(body=irs_matvec.uiet.body), - fielddata) + arrays) objs = self.objs cb = PETScCallable( @@ -125,7 +126,7 @@ def _make_matvec(self, fielddata, matvecs, prefix='MatMult'): self._matvecs.append(cb) self._efuncs[cb.name] = cb - def _create_matvec_body(self, body, fielddata): + def _create_matvec_body(self, body, arrays): linsolve_expr = self.injectsolve.expr.rhs objs = self.objs sobjs = self.solver_objs @@ -134,8 +135,8 @@ def _create_matvec_body(self, body, fielddata): ctx = objs['dummyctx'] xlocal = objs['xloc'] ylocal = objs['yloc'] - y_matvec = fielddata.arrays['y'] - x_matvec = fielddata.arrays['x'] + y_matvec = arrays['y'] + x_matvec = arrays['x'] body = self.timedep.uxreplace_time(body) @@ -652,18 +653,16 @@ def _make_core(self): all_fielddata = injectsolve.expr.rhs.fielddata for t in targets: - data = all_fielddata.get_field_data(t) - self._make_formfunc(data) - self._make_formrhs(data) - + row_matvecs = all_fielddata.submatrices.submatrices[t] + arrays = all_fielddata.arrays[t] for submat, mtvs in row_matvecs.items(): if mtvs['matvecs']: - self._make_matvec(data, mtvs['matvecs'], prefix=f'{submat}_MatMult') + self._make_matvec(arrays, mtvs['matvecs'], prefix=f'{submat}_MatMult') self._make_user_struct_callback() self._make_whole_matvec() - self._make_whole_formfunc() + self._make_whole_formfunc(all_fielddata) self._create_submatrices() self._efuncs['PopulateMatContext'] = self.objs['dummyefunc'] @@ -726,56 +725,175 @@ def _whole_matvec_body(self): retstmt=(Call('PetscFunctionReturn', arguments=[0]),) ) - def _make_whole_formfunc(self): + def _make_whole_formfunc(self, fielddata): + formfuncs = fielddata.formfuncs + # from IPython import embed; embed() + # Compile formfunc `eqns` into an IET via recursive compilation + irs_formfunc, _ = self.rcompile( + formfuncs, options={'mpi': False}, sregistry=self.sregistry, + concretize_mapper=self.concretize_mapper + ) + body_formfunc = self._whole_formfunc_body(List(body=irs_formfunc.uiet.body), + fielddata) objs = self.objs - body = self._whole_formfunc_body() - cb = PETScCallable( self.sregistry.make_name(prefix='WholeFormFunc'), - List(body=body), + body_formfunc, + # List(body=()), retval=objs['err'], parameters=(objs['snes'], objs['X'], objs['F'], objs['dummyptr']) ) self._main_formfunc_callback = cb self._efuncs[cb.name] = cb - def _whole_formfunc_body(self): + def _whole_formfunc_body(self, body, fielddata): + linsolve_expr = self.injectsolve.expr.rhs objs = self.objs sobjs = self.solver_objs - ljacctx = objs['ljacctx'] - struct_cast = DummyExpr(ljacctx, JacobianStructCast(objs['dummyptr'])) - X = objs['X'] - F = objs['F'] + dmda = sobjs['callbackdm'] + ctx = objs['dummyctx'] - targets = self.injectsolve.expr.rhs.fielddata.targets + body = self.timedep.uxreplace_time(body) - deref_subdms = Dereference(objs['LocalSubdms'], ljacctx) - deref_fields = Dereference(objs['LocalFields'], ljacctx) + fields = self._dummy_fields(body) + self._struct_params.extend(fields) - calls = () - for i, t in enumerate(targets): - field_ptr = FieldFromPointer(objs['LocalFields'].indexed[i], ljacctx) - x_name = f'Xglobal{t.name}' - f_name = f'Fglobal{t.name}' - calls += ( - petsc_call('VecGetSubVector', [X, field_ptr, Byref(sobjs[x_name])]), - petsc_call('VecGetSubVector', [F, field_ptr, Byref(sobjs[f_name])]), - petsc_call(self.formfuncs[i].name, [ - objs['snes'], sobjs[x_name], sobjs[f_name], - VOID(objs['LocalSubdms'].indexed[i], stars='*') - ]), - petsc_call('VecRestoreSubVector', [X, field_ptr, Byref(sobjs[x_name])]), - petsc_call('VecRestoreSubVector', [F, field_ptr, Byref(sobjs[f_name])]), - ) - return CallableBody( - List(body=calls + (BlankLine,)), - init=(objs['begin_user'],), - stacks=(struct_cast, deref_subdms, deref_fields), - retstmt=(Call('PetscFunctionReturn', arguments=[0]),) + targets = fielddata.targets + arrays = fielddata.arrays + + f_u = arrays[targets[0]]['f'] + f_v = arrays[targets[1]]['f'] + + x_u = arrays[targets[0]]['x'] + x_v = arrays[targets[1]]['x'] + + fbundle = PetscBundle(name='f_vu', components=(f_v, f_u)) + xbundle = PetscBundle(name='x_vu', components=(x_v, x_u)) + + mapper1 = {x_u.base: xbundle.base, x_v.base: xbundle.base} + + indexeds = FindSymbols('indexeds').visit(body) + + subss = {} + for i in indexeds: + if i.base in mapper1: + bundle = mapper1[i.base] + subss[i] = i.func(bundle, *i.indices) + # subss[i] = bundle.indexed.__getitem__((1,)+i.indices) + # subss[i] = subss[i].__getitem__((0,)+i.indices) + subss[i] = PetscComponentAccess(subss[i], 1) + + # subss[i] = xbundle.access_component() + + # from IPython import embed; embed() + + body = Uxreplace(subss).visit(body) + + dm_cast = DummyExpr(dmda, DMCast(objs['dummyptr']), init=True) + + dm_get_app_context = petsc_call( + 'DMGetApplicationContext', [dmda, Byref(ctx._C_symbol)] + ) + + dm_get_local_xvec = petsc_call( + 'DMGetLocalVector', [dmda, Byref(objs['xloc'])] + ) + + global_to_local_begin = petsc_call( + 'DMGlobalToLocalBegin', [dmda, objs['X'], + insert_vals, objs['xloc']] + ) + + global_to_local_end = petsc_call('DMGlobalToLocalEnd', [ + dmda, objs['X'], insert_vals, objs['xloc'] + ]) + + dm_get_local_yvec = petsc_call( + 'DMGetLocalVector', [dmda, Byref(objs['floc'])] + ) + + vec_get_array_f = petsc_call( + 'VecGetArray', [objs['floc'], Byref(fbundle.vector._C_symbol)] + ) + + vec_get_array_x = petsc_call( + 'VecGetArray', [objs['xloc'], Byref(xbundle.vector._C_symbol)] + ) + + dm_get_local_info = petsc_call( + 'DMDAGetLocalInfo', [dmda, Byref(linsolve_expr.localinfo)] + ) + + vec_restore_array_f = petsc_call( + 'VecRestoreArray', [objs['floc'], Byref(fbundle.vector._C_symbol)] + ) + + vec_restore_array_x = petsc_call( + 'VecRestoreArray', [objs['xloc'], Byref(xbundle.vector._C_symbol)] + ) + + dm_local_to_global_begin = petsc_call('DMLocalToGlobalBegin', [ + dmda, objs['floc'], insert_vals, objs['F'] + ]) + + dm_local_to_global_end = petsc_call('DMLocalToGlobalEnd', [ + dmda, objs['floc'], insert_vals, objs['F'] + ]) + + dm_restore_local_xvec = petsc_call( + 'DMRestoreLocalVector', [dmda, Byref(objs['xloc'])] + ) + + dm_restore_local_yvec = petsc_call( + 'DMRestoreLocalVector', [dmda, Byref(objs['floc'])] + ) + + body = body._rebuild( + body=body.body + + ( + vec_restore_array_f, + vec_restore_array_x, + dm_local_to_global_begin, + dm_local_to_global_end, + dm_restore_local_xvec, + dm_restore_local_yvec) + ) + + stacks = ( + # cast, + dm_cast, + dm_get_app_context, + dm_get_local_xvec, + global_to_local_begin, + global_to_local_end, + dm_get_local_yvec, + vec_get_array_f, + vec_get_array_x, + dm_get_local_info ) + # Dereference function data in struct + dereference_funcs = [Dereference(i, ctx) for i in + fields if isinstance(i.function, AbstractFunction)] + + + f_soa = PointerCast(fbundle) + x_soa = PointerCast(xbundle) + + formfunc_body = CallableBody( + List(body=body), + init=(objs['begin_user'],), + stacks=stacks+tuple(dereference_funcs), + casts=(f_soa,x_soa), + retstmt=(Call('PetscFunctionReturn', arguments=[0]),)) + + # Replace non-function data with pointer to data in struct + subs = {i._C_symbol: FieldFromPointer(i._C_symbol, ctx) for i in fields} + + return Uxreplace(subs).visit(formfunc_body) + def _create_submatrices(self): body = self._submat_callback_body() objs = self.objs @@ -1411,7 +1529,7 @@ def _setup(self): call_coupled_struct_callback, shell_set_ctx, create_submats) + \ - tuple(deref_dms) + tuple(xglobals) + tuple(bglobals) + tuple(deref_dms) + tuple(xglobals) return coupled_setup @@ -1501,7 +1619,7 @@ def _execute_solve(self): struct_assignment = self.timedep.assign_time_iters(sobjs['userctx']) - rhs_callbacks = self.cbbuilder.formrhs + # rhs_callbacks = self.cbbuilder.formrhs xglob = sobjs['xglobal'] bglob = sobjs['bglobal'] @@ -1512,7 +1630,7 @@ def _execute_solve(self): pre_solve = () post_solve = () - for i, (c, t) in enumerate(zip(rhs_callbacks, targets)): + for i, t in enumerate(targets): name = t.name dm = sobjs[f'da{name}'] target_xloc = sobjs[f'xlocal{name}'] @@ -1522,8 +1640,11 @@ def _execute_solve(self): s = sobjs[f'scatter{name}'] pre_solve += ( - petsc_call(c.name, [dm, target_bglob]), + # petsc_call(c.name, [dm, target_bglob]), + # TODO: switch to createwitharray and move to setup petsc_call('DMCreateLocalVector', [dm, Byref(target_xloc)]), + + # TODO: need to call reset array self.timedep.place_array(t), petsc_call( 'DMLocalToGlobal', @@ -1549,6 +1670,7 @@ def _execute_solve(self): 'VecScatterEnd', [s, target_bglob, bglob, insert_vals, sreverse] ), + BlankLine, ) post_solve += ( diff --git a/devito/petsc/solve.py b/devito/petsc/solve.py index a4917dfe8f..1a3e2bbf95 100644 --- a/devito/petsc/solve.py +++ b/devito/petsc/solve.py @@ -1,16 +1,17 @@ from functools import singledispatch import sympy +from collections import defaultdict from devito.finite_differences.differentiable import Mul from devito.finite_differences.derivative import Derivative from devito.types import Eq, Symbol, SteppingDimension, TimeFunction from devito.types.equation import PetscEq from devito.operations.solve import eval_time_derivatives -from devito.symbolics import retrieve_functions +from devito.symbolics import retrieve_functions, FieldFromComposite from devito.tools import as_tuple, filter_ordered from devito.petsc.types import (LinearSolveExpr, PETScArray, DMDALocalInfo, - FieldData, MultipleFieldData, SubMatrices) + FieldData, MultipleFieldData, SubMatrices, PetscBundle) __all__ = ['PETScSolve', 'EssentialBC'] @@ -73,14 +74,14 @@ def generate_field_data(self, eqns, target, arrays): ) def build_function_eqns(self, eq, target, arrays): - b, F_target, targets = separate_eqn(eq, target) + b, F_target, zeroed_eqn, targets = separate_eqn(eq, target) formfunc = self.make_formfunc(eq, F_target, arrays, targets) formrhs = self.make_rhs(eq, b, arrays) return (formfunc, formrhs) def build_matvec_eqns(self, eq, target, arrays): - b, F_target, targets = separate_eqn(eq, target) + b, F_target, zeroed_eqn, targets = separate_eqn(eq, target) if not F_target: return None matvec = self.make_matvec(eq, F_target, arrays, targets) @@ -98,8 +99,11 @@ def make_formfunc(self, eq, F_target, arrays, targets): if isinstance(eq, EssentialBC): rhs = 0. else: - rhs = F_target.subs(targets_to_arrays(arrays['x'], targets)) - rhs = rhs.subs(self.time_mapper) + if isinstance(F_target, (int, float)): + rhs = F_target + else: + rhs = F_target.subs(targets_to_arrays(arrays['x'], targets)) + rhs = rhs.subs(self.time_mapper) return Eq(arrays['f'], rhs, subdomain=eq.subdomain) def make_rhs(self, eq, b, arrays): @@ -140,21 +144,46 @@ def linear_solve_args(self): funcs = get_funcs(combined_eqns) self.time_mapper = generate_time_mapper(funcs) - targets = list(self.target_eqns.keys()) - jacobian = SubMatrices(targets) + coupled_targets = list(self.target_eqns.keys()) + jacobian = SubMatrices(coupled_targets) + + arrays = self.generate_arrays_combined(*coupled_targets) + + + # from IPython import embed; embed() + # TODO: don't need a 'b' for coupled + # arrays = self.generate_arrays_combined(*coupled_targets) + # all_formfuncs = [] + + # from IPython import embed; embed() + # f_field = [arrays[coupled_targets[t]]['f'] for t in range(len(coupled_targets))] + # f_field = PetscBundle(name='fuv_field', components=f_field) + + + all_data = MultipleFieldData(submatrices=jacobian, arrays=arrays, targets=coupled_targets) + + # f_v = arrays[coupled_targets[0]]['f'] + # f_u = arrays[coupled_targets[1]]['f'] + - all_data = MultipleFieldData(jacobian) for target, eqns in self.target_eqns.items(): eqns = as_tuple(eqns) - arrays = self.generate_arrays(target) - self.update_jacobian(eqns, target, jacobian, arrays) - fielddata = self.generate_field_data( - eqns, target, arrays - ) - all_data.add_field_data(fielddata) + # TODO: obvs fix and don't duplicate arrays here + # tmp_arr = self.generate_arrays(target) + self.update_jacobian(eqns, target, jacobian, arrays[target]) + + formfuncs = [self.build_function_eqns(eq, target, coupled_targets, arrays) for eq in eqns] + + all_data.extend_formfuncs(formfuncs) + + # all_formfuncs.extend(formfuncs) + + # all_data. + + # from IPython import embed; embed() return target, tuple(funcs), all_data @@ -164,21 +193,75 @@ def update_jacobian(self, eqns, target, jacobian, arrays): self.build_matvec_eqns(eq, mtvs['derivative_wrt'], arrays) for eq in eqns ] - # Set submatrix only if there's at least one non-zero matvec - if any(m is not None for m in matvecs): + matvecs = [m for m in matvecs if m is not None] + if matvecs: jacobian.set_submatrix(target, submat, matvecs) - def generate_field_data(self, eqns, target, arrays): - formfuncs, formrhs = zip( - *[self.build_function_eqns(eq, target, arrays) for eq in eqns] - ) + # def generate_field_data(self, eqns, target, arrays): + # # from IPython import embed; embed() + # formfuncs, formrhs = zip( + # *[self.build_function_eqns(eq, target, arrays) for eq in eqns] + # ) + + # return FieldData( + # target=target, + # formfuncs=formfuncs, + # formrhs=formrhs, + # arrays=arrays + # ) + + def build_function_eqns(self, eq, main_target, coupled_targets, arrays): + zeroed = eq.lhs - eq.rhs + + zeroed_eqn = Eq(zeroed, 0) + zeroed_eqn = eval_time_derivatives(zeroed) + + mapper = {} + for t in coupled_targets: + target_funcs = generate_targets(Eq(zeroed, 0), t) + mapper.update(targets_to_arrays(arrays[t]['x'], target_funcs)) + + # from IPython import embed; embed() + + if isinstance(eq, EssentialBC): + rhs = 0. + else: + if isinstance(zeroed, (int, float)): + rhs = zeroed + else: + rhs = zeroed.subs(mapper) + rhs = rhs.subs(self.time_mapper) + + return Eq(arrays[main_target]['f'], rhs, subdomain=eq.subdomain) + # return Eq(f_field, rhs, subdomain=eq.subdomain) + + def generate_arrays_combined(self, *targets): + return { + target: { + p: PETScArray( + name=f'{p}_{target.name}', + target=target, + liveness='eager', + localinfo=localinfo + ) + for p in prefixes + } + for target in targets + } + + + # def make_formfunc(self, eq, F_target, arrays, targets): + # if isinstance(eq, EssentialBC): + # rhs = 0. + # else: + # if isinstance(F_target, (int, float)): + # rhs = F_target + # else: + # # from IPython import embed; embed() + # rhs = F_target.subs(targets_to_arrays(arrays['x'], targets)) + # rhs = rhs.subs(self.time_mapper) + # return Eq(arrays['f'], rhs, subdomain=eq.subdomain) - return FieldData( - target=target, - formfuncs=formfuncs, - formrhs=formrhs, - arrays=arrays - ) class EssentialBC(Eq): @@ -194,7 +277,7 @@ def separate_eqn(eqn, target): zeroed_eqn = eval_time_derivatives(zeroed_eqn.lhs) target_funcs = set(generate_targets(zeroed_eqn, target)) b, F_target = remove_targets(zeroed_eqn, target_funcs) - return -b, F_target, target_funcs + return -b, F_target, zeroed_eqn, target_funcs def generate_targets(eq, target): diff --git a/devito/petsc/types/array.py b/devito/petsc/types/array.py index 54bbb4a9f4..3e35e82b2b 100644 --- a/devito/petsc/types/array.py +++ b/devito/petsc/types/array.py @@ -1,12 +1,13 @@ from functools import cached_property -from ctypes import POINTER +from ctypes import POINTER, Structure from devito.types.utils import DimensionTuple -from devito.types.array import ArrayBasic +from devito.types.array import ArrayBasic, Bundle, ArrayMapped, ComponentAccess from devito.finite_differences import Differentiable -from devito.types.basic import AbstractFunction -from devito.tools import dtype_to_ctype, as_tuple +from devito.types.basic import AbstractFunction, IndexedData +from devito.tools import dtype_to_ctype, as_tuple, dtypes_vector_mapper, CustomDtype from devito.symbolics import FieldFromComposite +from devito.petsc.types.object import PETScStruct class PETScArray(ArrayBasic, Differentiable): @@ -116,3 +117,74 @@ def symbolic_shape(self): FieldFromComposite('g%sm' % d.name, self.localinfo) for d in self.dimensions] # Reverse it since DMDA is setup backwards to Devito dimensions. return DimensionTuple(*field_from_composites[::-1], getters=self.dimensions) + + +class PetscBundle(Bundle): + """ + """ + + is_Bundle = True + _data_alignment = False + + @property + def _C_ctype(self): + # TODO: extend to cases with multiple petsc solves... + fields = [(i.name, dtype_to_ctype(i.dtype)) for i in self.components] + return POINTER(type('Field', (Structure,), {'_fields_': fields})) + + @cached_property + def indexed(self): + """The wrapped IndexedData object.""" + return AoSIndexedData(self.name, shape=self._shape, function=self.function) + + @cached_property + def vector(self): + return PETScArray( + name=self.name, + target=self.c0.target, + liveness=self.c0.liveness, + localinfo=self.c0.localinfo, + ) + + @property + def _C_name(self): + return self.vector._C_name + + def __getitem__(self, index): + index = as_tuple(index) + if len(index) == self.ndim: + return super().__getitem__(index) + elif len(index) == self.ndim + 1: + component_index, indices = index[0], index[1:] + # from IPython import embed; embed() + return ComponentAccess(self.indexed[indices], component_index) + else: + raise ValueError("Expected %d or %d indices, got %d instead" + % (self.ndim, self.ndim + 1, len(index))) + + # def access_component(self): + # _component_names = tuple(i.name for i in components) + # class PetscComponentAccess(ComponentAccess): + # _component_names = _component_names + + # # components = self.components + + # # _component_names = tuple(i.name for i in components) + # from IPython import embed; embed() + + # return component + + + +class AoSIndexedData(IndexedData): + @property + def dtype(self): + return self.function._C_ctype + + + + +class PetscComponentAccess(ComponentAccess): + + # _component_names = ('x_u', 'x_v') + pass diff --git a/devito/petsc/types/object.py b/devito/petsc/types/object.py index 2e69e4a0fc..62bccc5b2c 100644 --- a/devito/petsc/types/object.py +++ b/devito/petsc/types/object.py @@ -308,3 +308,10 @@ class ArgvSymbol(DataSymbol): @property def _C_ctype(self): return POINTER(POINTER(c_char)) + + + +# class PetscArrayofStruct(LocalCompositeObject): +# pass +# """ +# """ \ No newline at end of file diff --git a/devito/petsc/types/types.py b/devito/petsc/types/types.py index bfbdd4972b..25da0d5575 100644 --- a/devito/petsc/types/types.py +++ b/devito/petsc/types/types.py @@ -1,8 +1,10 @@ import sympy -from devito.tools import Reconstructable, sympy_mutex +from devito.tools import Reconstructable, sympy_mutex, as_tuple from devito.tools.dtypes_lowering import dtype_mapper from devito.petsc.utils import petsc_variables +from devito.petsc.types.object import PETScStruct +# from devito.petsc.types.array import PetscArrayofStruct class MetaData(sympy.Function, Reconstructable): @@ -133,7 +135,7 @@ def eval(cls, *args): class FieldData: def __init__(self, target=None, matvecs=None, formfuncs=None, formrhs=None, initialguess=None, arrays=None, **kwargs): - self._target = kwargs.get('target', target) + self._target = kwargs.get('target', None) petsc_precision = dtype_mapper[petsc_variables['PETSC_PRECISION']] if self._target.dtype != petsc_precision: @@ -189,32 +191,42 @@ def targets(self): return (self.target,) +# TODO: should this acc inherhit from fielddata? maybe not? class MultipleFieldData(FieldData): - def __init__(self, submatrices=None): - self.field_data_list = [] + def __init__(self, targets, arrays, submatrices=None): + self._targets = as_tuple(targets) + self._arrays = arrays + # self.field_data_list = [] self._submatrices = submatrices + self._formfuncs = [] + # self._f_field = f_field - def add_field_data(self, field_data): - self.field_data_list.append(field_data) + # def add_field_data(self, field_data): + # self.field_data_list.append(field_data) - def get_field_data(self, target): - for field_data in self.field_data_list: - if field_data.target == target: - return field_data - raise ValueError(f"FieldData with target {target} not found.") + # def get_field_data(self, target): + # for field_data in self.field_data_list: + # if field_data.target == target: + # return field_data + # raise ValueError(f"FieldData with target {target} not found.") pass - @property - def target(self): - return None + def extend_formfuncs(self, formfuncs): + self._formfuncs.extend(formfuncs) - @property - def targets(self): - return tuple(field_data.target for field_data in self.field_data_list) + # @property + # def target(self): + # return None + + # @property + # def targets(self): + # return tuple(field_data.target for field_data in self.field_data_list) @property def space_dimensions(self): - space_dims = {field_data.space_dimensions for field_data in self.field_data_list} + # space_dims = {field_data.space_dimensions for field_data in self.field_data_list} + space_dims = {t.space_dimensions for t in self.targets} + # from IPython import embed; embed() if len(space_dims) > 1: # TODO: This may not actually have to be the case, but enforcing it for now raise ValueError( @@ -249,6 +261,18 @@ def space_order(self): def submatrices(self): return self._submatrices + # @property + # def formfuncs(self): + # return self._formfuncs + + @property + def targets(self): + return self._targets + + @property + def arrays(self): + return self._arrays + class SubMatrices: def __init__(self, targets): diff --git a/devito/types/basic.py b/devito/types/basic.py index a28ca2f486..9990c55836 100644 --- a/devito/types/basic.py +++ b/devito/types/basic.py @@ -87,6 +87,7 @@ def _C_typedata(self): if isinstance(_type, CustomDtype): return _type + # from IPython import embed; embed() while issubclass(_type, _Pointer): _type = _type._type_ From f6b38eddb4c3c7c85c93854bcbc76ff41ed1942d Mon Sep 17 00:00:00 2001 From: ZoeLeibowitz Date: Tue, 29 Apr 2025 17:40:00 +0100 Subject: [PATCH 26/84] compiler: Working petscbundle --- devito/petsc/iet/routines.py | 44 ++++++++++++++++-------------------- devito/petsc/types/array.py | 35 ++++++++-------------------- devito/petsc/types/types.py | 3 +-- devito/types/array.py | 8 +++++-- 4 files changed, 36 insertions(+), 54 deletions(-) diff --git a/devito/petsc/iet/routines.py b/devito/petsc/iet/routines.py index be76ceace1..3daa7efcfc 100644 --- a/devito/petsc/iet/routines.py +++ b/devito/petsc/iet/routines.py @@ -21,7 +21,7 @@ from devito.petsc.types import (DM, Mat, CallbackVec, Vec, KSP, PC, SNES, PetscInt, StartPtr, PointerIS, PointerDM, VecScatter, DMCast, JacobianStructCast, JacobianStruct, - SubMatrixStruct, CallbackDM, PetscComponentAccess) + SubMatrixStruct, CallbackDM) class CBBuilder: @@ -660,9 +660,9 @@ def _make_core(self): if mtvs['matvecs']: self._make_matvec(arrays, mtvs['matvecs'], prefix=f'{submat}_MatMult') - self._make_user_struct_callback() self._make_whole_matvec() self._make_whole_formfunc(all_fielddata) + self._make_user_struct_callback() self._create_submatrices() self._efuncs['PopulateMatContext'] = self.objs['dummyefunc'] @@ -758,7 +758,7 @@ def _whole_formfunc_body(self, body, fielddata): fields = self._dummy_fields(body) self._struct_params.extend(fields) - + # from IPython import embed; embed() targets = fielddata.targets arrays = fielddata.arrays @@ -769,10 +769,16 @@ def _whole_formfunc_body(self, body, fielddata): x_u = arrays[targets[0]]['x'] x_v = arrays[targets[1]]['x'] - fbundle = PetscBundle(name='f_vu', components=(f_v, f_u)) - xbundle = PetscBundle(name='x_vu', components=(x_v, x_u)) - mapper1 = {x_u.base: xbundle.base, x_v.base: xbundle.base} + target_indices = {t: i for i, t in enumerate(targets)} + # from IPython import embed; embed() + + # TODO: to group them, maybe pass in struct name as arg to petscbundle + fbundle = PetscBundle(name='f_vu', components=(f_u, f_v)) + xbundle = PetscBundle(name='x_vu', components=(x_u, x_v)) + + mapper1 = {x_u.base: xbundle, x_v.base: xbundle, f_u.base: fbundle, + f_v.base: fbundle} indexeds = FindSymbols('indexeds').visit(body) @@ -780,14 +786,10 @@ def _whole_formfunc_body(self, body, fielddata): for i in indexeds: if i.base in mapper1: bundle = mapper1[i.base] - subss[i] = i.func(bundle, *i.indices) - # subss[i] = bundle.indexed.__getitem__((1,)+i.indices) - # subss[i] = subss[i].__getitem__((0,)+i.indices) - subss[i] = PetscComponentAccess(subss[i], 1) - - # subss[i] = xbundle.access_component() - - # from IPython import embed; embed() + index = target_indices[i.function.target] + index = (index,)+i.indices + subss[i] = bundle.__getitem__(index) + body = Uxreplace(subss).visit(body) @@ -1373,9 +1375,9 @@ def _create_dmda(self, dmda): class CoupledSetup(BaseSetup): - @property - def snes_ctx(self): - return Byref(self.solver_objs['jacctx']) + # @property + # def snes_ctx(self): + # return Byref(self.solver_objs['jacctx']) def _setup(self): # TODO: minimise code duplication with superclass @@ -1662,14 +1664,6 @@ def _execute_solve(self): 'VecScatterEnd', [s, target_xglob, xglob, insert_vals, sreverse] ), - petsc_call( - 'VecScatterBegin', - [s, target_bglob, bglob, insert_vals, sreverse] - ), - petsc_call( - 'VecScatterEnd', - [s, target_bglob, bglob, insert_vals, sreverse] - ), BlankLine, ) diff --git a/devito/petsc/types/array.py b/devito/petsc/types/array.py index 3e35e82b2b..39eb2e2347 100644 --- a/devito/petsc/types/array.py +++ b/devito/petsc/types/array.py @@ -1,6 +1,8 @@ from functools import cached_property from ctypes import POINTER, Structure +from sympy import Expr + from devito.types.utils import DimensionTuple from devito.types.array import ArrayBasic, Bundle, ArrayMapped, ComponentAccess from devito.finite_differences import Differentiable @@ -122,14 +124,14 @@ def symbolic_shape(self): class PetscBundle(Bundle): """ """ - is_Bundle = True + _data_alignment = False @property def _C_ctype(self): # TODO: extend to cases with multiple petsc solves... - fields = [(i.name, dtype_to_ctype(i.dtype)) for i in self.components] + fields = [(i.target.name, dtype_to_ctype(i.dtype)) for i in self.components] return POINTER(type('Field', (Structure,), {'_fields_': fields})) @cached_property @@ -156,35 +158,18 @@ def __getitem__(self, index): return super().__getitem__(index) elif len(index) == self.ndim + 1: component_index, indices = index[0], index[1:] - # from IPython import embed; embed() - return ComponentAccess(self.indexed[indices], component_index) + names = tuple(i.target.name for i in self.components) + return ComponentAccess( + self.indexed[indices], + component_index, + component_names=names + ) else: raise ValueError("Expected %d or %d indices, got %d instead" % (self.ndim, self.ndim + 1, len(index))) - # def access_component(self): - # _component_names = tuple(i.name for i in components) - # class PetscComponentAccess(ComponentAccess): - # _component_names = _component_names - - # # components = self.components - - # # _component_names = tuple(i.name for i in components) - # from IPython import embed; embed() - - # return component - - class AoSIndexedData(IndexedData): @property def dtype(self): return self.function._C_ctype - - - - -class PetscComponentAccess(ComponentAccess): - - # _component_names = ('x_u', 'x_v') - pass diff --git a/devito/petsc/types/types.py b/devito/petsc/types/types.py index 25da0d5575..e470986612 100644 --- a/devito/petsc/types/types.py +++ b/devito/petsc/types/types.py @@ -135,8 +135,7 @@ def eval(cls, *args): class FieldData: def __init__(self, target=None, matvecs=None, formfuncs=None, formrhs=None, initialguess=None, arrays=None, **kwargs): - self._target = kwargs.get('target', None) - + self._target = target petsc_precision = dtype_mapper[petsc_variables['PETSC_PRECISION']] if self._target.dtype != petsc_precision: raise TypeError( diff --git a/devito/types/array.py b/devito/types/array.py index eb1fd0dc4c..32c8ccc97b 100644 --- a/devito/types/array.py +++ b/devito/types/array.py @@ -495,6 +495,7 @@ def _C_get_field(self, region, dim, side=None): return self.c0._C_get_field(region, dim, side=side) def __getitem__(self, index): + # from IPython import embed; embed() index = as_tuple(index) if len(index) == self.ndim: return super().__getitem__(index) @@ -535,19 +536,22 @@ def handles(self): class ComponentAccess(Expr, Pickable): - _component_names = ('x', 'y', 'z', 'w') + _default_component_names = ('x', 'y', 'z', 'w') __rargs__ = ('arg',) __rkwargs__ = ('index',) - def __new__(cls, arg, index=0, **kwargs): + def __new__(cls, arg, index=0, component_names=None, **kwargs): if not arg.is_Indexed: raise ValueError("Expected Indexed, got `%s` instead" % type(arg)) if not is_integer(index) or index > 3: raise ValueError("Expected 0 <= index < 4") + names = component_names or cls._default_component_names + obj = Expr.__new__(cls, arg) obj._index = index + obj._component_names = names return obj From 7300c8f6a896f86d751090a2ed4e7ec13c9ec530 Mon Sep 17 00:00:00 2001 From: ZoeLeibowitz Date: Wed, 30 Apr 2025 18:52:02 +0100 Subject: [PATCH 27/84] compiler: Add PetscMixin to simplify priority destroys --- .github/workflows/pytest-petsc.yml | 1 + devito/passes/iet/definitions.py | 2 +- devito/passes/iet/engine.py | 33 +- devito/passes/iet/misc.py | 3 +- devito/petsc/iet/routines.py | 181 ++++-- devito/petsc/solve.py | 82 +-- devito/petsc/types/array.py | 20 +- devito/petsc/types/object.py | 21 +- devito/petsc/types/types.py | 28 +- devito/tools/utils.py | 1 + devito/types/array.py | 1 - devito/types/basic.py | 1 - .../petsc/random/biharmonic/02_biharmonic.py | 148 +++++ .../random/biharmonic/biharmonic_matfree.c | 556 ++++++++++++++++++ .../biharmonic/biharmonic_matfree_nonscaled.c | 553 +++++++++++++++++ tests/test_petsc.py | 4 + 16 files changed, 1440 insertions(+), 195 deletions(-) create mode 100644 examples/petsc/random/biharmonic/02_biharmonic.py create mode 100644 examples/petsc/random/biharmonic/biharmonic_matfree.c create mode 100644 examples/petsc/random/biharmonic/biharmonic_matfree_nonscaled.c diff --git a/.github/workflows/pytest-petsc.yml b/.github/workflows/pytest-petsc.yml index 13177cd7b8..19446642af 100644 --- a/.github/workflows/pytest-petsc.yml +++ b/.github/workflows/pytest-petsc.yml @@ -83,6 +83,7 @@ jobs: ${{ env.RUN_CMD }} mpiexec -n 1 python3 examples/petsc/Poisson/03_poisson.py ${{ env.RUN_CMD }} mpiexec -n 1 python3 examples/petsc/Poisson/04_poisson.py ${{ env.RUN_CMD }} mpiexec -n 1 python3 examples/petsc/random/01_helmholtz.py + ${{ env.RUN_CMD }} mpiexec -n 1 python3 examples/petsc/random/biharmonic/02_biharmonic.py - name: Upload coverage to Codecov if: "!contains(matrix.name, 'docker')" diff --git a/devito/passes/iet/definitions.py b/devito/passes/iet/definitions.py index f5ae337b9c..e56c86645f 100644 --- a/devito/passes/iet/definitions.py +++ b/devito/passes/iet/definitions.py @@ -379,7 +379,7 @@ def place_definitions(self, iet, globs=None, **kwargs): # created by the compiler up to this point (Array, LocalObject, etc.) storage = Storage() defines = FindSymbols('defines-aliases|globals').visit(iet) - # from IPython import embed; embed() + for i in FindSymbols().visit(iet): if i in defines: continue diff --git a/devito/passes/iet/engine.py b/devito/passes/iet/engine.py index 992096fd41..a90f4ba275 100644 --- a/devito/passes/iet/engine.py +++ b/devito/passes/iet/engine.py @@ -13,9 +13,11 @@ from devito.types import (Array, Bundle, CompositeObject, Lock, IncrDimension, ModuloDimension, Indirection, Pointer, SharedData, ThreadArray, Temp, NPThreads, NThreadsBase, Wildcard) +from devito.types.array import ArrayBasic from devito.types.args import ArgProvider from devito.types.dense import DiscreteFunction from devito.types.dimension import AbstractIncrDimension, BlockDimension +from devito.petsc.types import PETScArray __all__ = ['Graph', 'iet_pass', 'iet_visit'] @@ -149,8 +151,6 @@ def apply(self, func, **kwargs): if len(efuncs) > len(self.efuncs): efuncs = reuse_compounds(efuncs, self.sregistry) # TODO: fix for petsc bundles - # TODO: somethinng to do with this bug is causing the compiler not to combine loops with petsc bundles - # from IPython import embed; embed() # efuncs = reuse_efuncs(self.root, efuncs, self.sregistry) self.efuncs = efuncs @@ -355,7 +355,7 @@ def reuse_efuncs(root, efuncs, sregistry=None): if isinstance(efunc, AsyncCallable): mapper[len(mapper)] = (efunc, [efunc]) continue - + # from IPython import embed; embed() afunc = abstract_efunc(efunc) key = afunc._signature() @@ -394,7 +394,7 @@ def abstract_efunc(efunc): - Objects are renamed as "o0", "o1", ... """ functions = FindSymbols('basics|symbolics|dimensions').visit(efunc) - # from IPython import embed; embed() + mapper = abstract_objects(functions) # from IPython import embed; embed() efunc = Uxreplace(mapper).visit(efunc) @@ -411,24 +411,24 @@ def abstract_objects(objects0, sregistry=None): objects = [] for i in objects0: if i.is_Bundle: - # from IPython import embed; embed() objects.extend(i.components) objects.append(i) - # from IPython import embed; embed() + # Precedence rules make it possible to reconstruct objects that depend on # higher priority objects - keys = [Bundle, Array, DiscreteFunction, AbstractIncrDimension, BlockDimension] + # keys = [Bundle, Array, PETScArray, DiscreteFunction, AbstractIncrDimension, BlockDimension] + keys = [Bundle, PETScArray, DiscreteFunction, AbstractIncrDimension, BlockDimension] priority = {k: i for i, k in enumerate(keys, start=1)} objects = sorted_priority(objects, priority) # Build abstraction mappings mapper = {} sregistry = sregistry or SymbolRegistry() - # from IPython import embed; embed() + for i in objects: - # from IPython import embed; embed() abstract_object(i, mapper, sregistry) + # from IPython import embed; embed() return mapper @@ -472,9 +472,24 @@ def _(i, mapper, sregistry): mapper[i.dmap] = v.dmap +# @abstract_object.register(PETScArray) +# def _(i, mapper, sregistry): +# name = sregistry.make_name(prefix='xx') + +# v = i._rebuild(name=name, initializer=None, alias=True) + +# mapper.update({ +# i: v, +# i.indexed: v.indexed, +# i.dmap: v.dmap, +# i._C_symbol: v._C_symbol, +# }) + + @abstract_object.register(Bundle) def _(i, mapper, sregistry): name = sregistry.make_name(prefix='a') + components = [mapper[f] for f in i.components] v = i._rebuild(name=name, components=components, alias=True) diff --git a/devito/passes/iet/misc.py b/devito/passes/iet/misc.py index 0f15deec81..e404a8e373 100644 --- a/devito/passes/iet/misc.py +++ b/devito/passes/iet/misc.py @@ -184,7 +184,6 @@ def _generate_macros_findexeds(iet, sregistry=None, tracker=None, **kwargs): try: v = tracker[i.base].v subs[i] = v.func(v.base, *i.indices) - # from IPython import embed; embed() continue except KeyError: pass @@ -194,7 +193,7 @@ def _generate_macros_findexeds(iet, sregistry=None, tracker=None, **kwargs): subs[i] = v tracker[i.base] = Bunch(header=header, v=v) - # from IPython import embed; embed() + iet = Uxreplace(subs).visit(iet) return iet diff --git a/devito/petsc/iet/routines.py b/devito/petsc/iet/routines.py index 3daa7efcfc..8674dc4df4 100644 --- a/devito/petsc/iet/routines.py +++ b/devito/petsc/iet/routines.py @@ -5,9 +5,9 @@ from devito.ir.iet import (Call, FindSymbols, List, Uxreplace, CallableBody, Dereference, DummyExpr, BlankLine, Callable, FindNodes, retrieve_iteration_tree, filter_iterations, Iteration, - PointerCast, Expression, Transformer) + PointerCast) from devito.symbolics import (Byref, FieldFromPointer, cast, VOID, - FieldFromComposite, IntDiv, Deref, Mod, uxreplace) + FieldFromComposite, IntDiv, Deref, Mod) from devito.symbolics.unevaluation import Mul from devito.types.basic import AbstractFunction from devito.types import Temp, Dimension @@ -20,7 +20,7 @@ from devito.petsc.utils import solver_mapper from devito.petsc.types import (DM, Mat, CallbackVec, Vec, KSP, PC, SNES, PetscInt, StartPtr, PointerIS, PointerDM, VecScatter, - DMCast, JacobianStructCast, JacobianStruct, + DMCast, JacobianStruct, SubMatrixStruct, CallbackDM) @@ -99,6 +99,13 @@ def initialguesses(self): def user_struct_callback(self): return self._user_struct_callback + @property + def zero_memory(self): + """Indicates whether the memory of the output + vector should be set to zero before the computation + in the callback.""" + return True + def _make_core(self): fielddata = self.injectsolve.expr.rhs.fielddata self._make_matvec(fielddata.arrays, fielddata.matvecs) @@ -148,6 +155,10 @@ def _create_matvec_body(self, body, arrays): 'DMGetApplicationContext', [dmda, Byref(ctx._C_symbol)] ) + zero_y_memory = petsc_call( + 'VecSet', [objs['Y'], 0.0] + ) if self.zero_memory else None + dm_get_local_xvec = petsc_call( 'DMGetLocalVector', [dmda, Byref(xlocal)] ) @@ -165,6 +176,10 @@ def _create_matvec_body(self, body, arrays): 'DMGetLocalVector', [dmda, Byref(ylocal)] ) + zero_ylocal_memory = petsc_call( + 'VecSet', [ylocal, 0.0] + ) + vec_get_array_y = petsc_call( 'VecGetArray', [ylocal, Byref(y_matvec._C_symbol)] ) @@ -186,11 +201,11 @@ def _create_matvec_body(self, body, arrays): ) dm_local_to_global_begin = petsc_call('DMLocalToGlobalBegin', [ - dmda, ylocal, insert_vals, objs['Y'] + dmda, ylocal, add_vals, objs['Y'] ]) dm_local_to_global_end = petsc_call('DMLocalToGlobalEnd', [ - dmda, ylocal, insert_vals, objs['Y'] + dmda, ylocal, add_vals, objs['Y'] ]) dm_restore_local_xvec = petsc_call( @@ -221,10 +236,12 @@ def _create_matvec_body(self, body, arrays): stacks = ( mat_get_dm, dm_get_app_context, + zero_y_memory, dm_get_local_xvec, global_to_local_begin, global_to_local_end, dm_get_local_yvec, + zero_ylocal_memory, vec_get_array_y, vec_get_array_x, dm_get_local_info @@ -290,6 +307,10 @@ def _create_formfunc_body(self, body, fielddata): 'DMGetApplicationContext', [dmda, Byref(ctx._C_symbol)] ) + zero_f_memory = petsc_call( + 'VecSet', [objs['F'], 0.0] + ) if self.zero_memory else None + dm_get_local_xvec = petsc_call( 'DMGetLocalVector', [dmda, Byref(objs['xloc'])] ) @@ -328,11 +349,11 @@ def _create_formfunc_body(self, body, fielddata): ) dm_local_to_global_begin = petsc_call('DMLocalToGlobalBegin', [ - dmda, objs['floc'], insert_vals, objs['F'] + dmda, objs['floc'], add_vals, objs['F'] ]) dm_local_to_global_end = petsc_call('DMLocalToGlobalEnd', [ - dmda, objs['floc'], insert_vals, objs['F'] + dmda, objs['floc'], add_vals, objs['F'] ]) dm_restore_local_xvec = petsc_call( @@ -356,6 +377,7 @@ def _create_formfunc_body(self, body, fielddata): stacks = ( dm_cast, dm_get_app_context, + zero_f_memory, dm_get_local_xvec, global_to_local_begin, global_to_local_end, @@ -647,13 +669,19 @@ def main_matvec_callback(self): def main_formfunc_callback(self): return self._main_formfunc_callback + @property + def zero_memory(self): + """Indicates whether the memory of the output + vector should be set to zero before the computation + in the callback.""" + return False + def _make_core(self): injectsolve = self.injectsolve targets = injectsolve.expr.rhs.fielddata.targets all_fielddata = injectsolve.expr.rhs.fielddata for t in targets: - row_matvecs = all_fielddata.submatrices.submatrices[t] arrays = all_fielddata.arrays[t] for submat, mtvs in row_matvecs.items(): @@ -688,6 +716,10 @@ def _whole_matvec_body(self): nonzero_submats = self.submatrices.nonzero_submatrix_keys + zero_y_memory = petsc_call( + 'VecSet', [objs['Y'], 0.0] + ) + calls = () for sm in nonzero_submats: idx = self.submatrices.submat_to_index[sm] @@ -720,14 +752,13 @@ def _whole_matvec_body(self): ), ) return CallableBody( - List(body=(ctx_main, BlankLine) + calls), + List(body=(ctx_main, zero_y_memory, BlankLine) + calls), init=(objs['begin_user'],), retstmt=(Call('PetscFunctionReturn', arguments=[0]),) ) def _make_whole_formfunc(self, fielddata): formfuncs = fielddata.formfuncs - # from IPython import embed; embed() # Compile formfunc `eqns` into an IET via recursive compilation irs_formfunc, _ = self.rcompile( formfuncs, options={'mpi': False}, sregistry=self.sregistry, @@ -739,7 +770,6 @@ def _make_whole_formfunc(self, fielddata): cb = PETScCallable( self.sregistry.make_name(prefix='WholeFormFunc'), body_formfunc, - # List(body=()), retval=objs['err'], parameters=(objs['snes'], objs['X'], objs['F'], objs['dummyptr']) ) @@ -758,40 +788,12 @@ def _whole_formfunc_body(self, body, fielddata): fields = self._dummy_fields(body) self._struct_params.extend(fields) - # from IPython import embed; embed() - - targets = fielddata.targets - arrays = fielddata.arrays - - f_u = arrays[targets[0]]['f'] - f_v = arrays[targets[1]]['f'] - - x_u = arrays[targets[0]]['x'] - x_v = arrays[targets[1]]['x'] - - - target_indices = {t: i for i, t in enumerate(targets)} - # from IPython import embed; embed() - - # TODO: to group them, maybe pass in struct name as arg to petscbundle - fbundle = PetscBundle(name='f_vu', components=(f_u, f_v)) - xbundle = PetscBundle(name='x_vu', components=(x_u, x_v)) - mapper1 = {x_u.base: xbundle, x_v.base: xbundle, f_u.base: fbundle, - f_v.base: fbundle} - - indexeds = FindSymbols('indexeds').visit(body) - - subss = {} - for i in indexeds: - if i.base in mapper1: - bundle = mapper1[i.base] - index = target_indices[i.function.target] - index = (index,)+i.indices - subss[i] = bundle.__getitem__(index) - - - body = Uxreplace(subss).visit(body) + # Process body for residual callback, including generating bundles etc + bundles = sobjs['bundles'] + fbundle = bundles['f'] + xbundle = bundles['x'] + body = self.bundle_residual(body, bundles) dm_cast = DummyExpr(dmda, DMCast(objs['dummyptr']), init=True) @@ -799,6 +801,10 @@ def _whole_formfunc_body(self, body, fielddata): 'DMGetApplicationContext', [dmda, Byref(ctx._C_symbol)] ) + zero_f_memory = petsc_call( + 'VecSet', [objs['F'], 0.0] + ) + dm_get_local_xvec = petsc_call( 'DMGetLocalVector', [dmda, Byref(objs['xloc'])] ) @@ -837,11 +843,11 @@ def _whole_formfunc_body(self, body, fielddata): ) dm_local_to_global_begin = petsc_call('DMLocalToGlobalBegin', [ - dmda, objs['floc'], insert_vals, objs['F'] + dmda, objs['floc'], add_vals, objs['F'] ]) dm_local_to_global_end = petsc_call('DMLocalToGlobalEnd', [ - dmda, objs['floc'], insert_vals, objs['F'] + dmda, objs['floc'], add_vals, objs['F'] ]) dm_restore_local_xvec = petsc_call( @@ -854,8 +860,7 @@ def _whole_formfunc_body(self, body, fielddata): body = body._rebuild( body=body.body + - ( - vec_restore_array_f, + (vec_restore_array_f, vec_restore_array_x, dm_local_to_global_begin, dm_local_to_global_end, @@ -864,9 +869,9 @@ def _whole_formfunc_body(self, body, fielddata): ) stacks = ( - # cast, dm_cast, dm_get_app_context, + zero_f_memory, dm_get_local_xvec, global_to_local_begin, global_to_local_end, @@ -880,7 +885,6 @@ def _whole_formfunc_body(self, body, fielddata): dereference_funcs = [Dereference(i, ctx) for i in fields if isinstance(i.function, AbstractFunction)] - f_soa = PointerCast(fbundle) x_soa = PointerCast(xbundle) @@ -888,8 +892,9 @@ def _whole_formfunc_body(self, body, fielddata): List(body=body), init=(objs['begin_user'],), stacks=stacks+tuple(dereference_funcs), - casts=(f_soa,x_soa), - retstmt=(Call('PetscFunctionReturn', arguments=[0]),)) + casts=(f_soa, x_soa), + retstmt=(Call('PetscFunctionReturn', arguments=[0]),) + ) # Replace non-function data with pointer to data in struct subs = {i._C_symbol: FieldFromPointer(i._C_symbol, ctx) for i in fields} @@ -1041,6 +1046,26 @@ def _submat_callback_body(self): stacks=(get_ctx, deref_subdm), retstmt=(Call('PetscFunctionReturn', arguments=[0]),) ) + + def bundle_residual(self, body, bundles): + fbundle = bundles['f'] + xbundle = bundles['x'] + + mapper1 = bundles['bundle_mapper'] + + indexeds = FindSymbols('indexeds').visit(body) + + subss = {} + for i in indexeds: + if i.base in mapper1: + bundle = mapper1[i.base] + index = bundles['target_indices'][i.function.target] + index = (index,)+i.indices + subss[i] = bundle.__getitem__(index) + + body = Uxreplace(subss).visit(body) + return body + class BaseObjectBuilder: @@ -1120,6 +1145,7 @@ def _extend_build(self, base_dict): sreg = self.sregistry objs = self.objs targets = self.fielddata.targets + arrays = self.fielddata.arrays base_dict['fields'] = PointerIS( name=sreg.make_name(prefix='fields'), nindices=len(targets) @@ -1154,6 +1180,40 @@ def _extend_build(self, base_dict): base_dict[f'{key}Y'] = CallbackVec(f'{key}Y') base_dict[f'{key}F'] = CallbackVec(f'{key}F') + + # Bundle objects/metadata required by the coupled residual callback + f_components = [] + x_components = [] + bundle_mapper = {} + + target_indices = {t: i for i, t in enumerate(targets)} + + for t in targets: + f_arr = arrays[t]['f'] + x_arr = arrays[t]['x'] + f_components.append(f_arr) + x_components.append(x_arr) + + # TODO: to group them, maybe pass in struct name as arg to petscbundle + fbundle = PetscBundle(name='f_bundle', components=f_components) + xbundle = PetscBundle(name='x_bundle', components=x_components) + + # Build the bundle mapper + for i, t in enumerate(targets): + f_arr = arrays[t]['f'] + x_arr = arrays[t]['x'] + bundle_mapper[f_arr.base] = fbundle + bundle_mapper[x_arr.base] = xbundle + + + base_dict['bundles'] = { + 'f': fbundle, + 'x': xbundle, + 'bundle_mapper': bundle_mapper, + # TODO: maybe this shouldn't be here + 'target_indices': target_indices + } + return base_dict def _target_dependent(self, base_dict): @@ -1410,9 +1470,6 @@ def _setup(self): get_local_size = petsc_call('VecGetSize', [sobjs['xlocal'], Byref(sobjs['localsize'])]) - global_b = petsc_call('DMCreateGlobalVector', - [dmda, Byref(sobjs['bglobal'])]) - snes_get_ksp = petsc_call('SNESGetKSP', [sobjs['snes'], Byref(sobjs['ksp'])]) @@ -1500,11 +1557,6 @@ def _setup(self): [sobjs[f'da{t.name}'], Byref(sobjs[f'xglobal{t.name}'])] ) for t in targets] - bglobals = [petsc_call( - 'DMCreateGlobalVector', - [sobjs[f'da{t.name}'], Byref(sobjs[f'bglobal{t.name}'])] - ) for t in targets] - coupled_setup = dmda_calls + ( snes_create, snes_set_dm, @@ -1514,7 +1566,6 @@ def _setup(self): global_x, local_x, get_local_size, - global_b, snes_get_ksp, ksp_set_tols, ksp_set_type, @@ -1617,6 +1668,7 @@ def _execute_solve(self): Assigns the required time iterators to the struct and executes the necessary calls to execute the SNES solver. """ + objs = self.objs sobjs = self.solver_objs struct_assignment = self.timedep.assign_time_iters(sobjs['userctx']) @@ -1637,13 +1689,11 @@ def _execute_solve(self): dm = sobjs[f'da{name}'] target_xloc = sobjs[f'xlocal{name}'] target_xglob = sobjs[f'xglobal{name}'] - target_bglob = sobjs[f'bglobal{name}'] field = sobjs['fields'].indexed[i] s = sobjs[f'scatter{name}'] pre_solve += ( - # petsc_call(c.name, [dm, target_bglob]), - # TODO: switch to createwitharray and move to setup + # TODO: switch to createwitharray and move to setup petsc_call('DMCreateLocalVector', [dm, Byref(target_xloc)]), # TODO: need to call reset array @@ -1682,7 +1732,7 @@ def _execute_solve(self): ) ) - snes_solve = (petsc_call('SNESSolve', [sobjs['snes'], bglob, xglob]),) + snes_solve = (petsc_call('SNESSolve', [sobjs['snes'], objs['Null'], xglob]),) return List( body=( @@ -1893,5 +1943,6 @@ def assign_time_iters(self, struct): void = 'void' insert_vals = 'INSERT_VALUES' +add_vals = 'ADD_VALUES' sreverse = 'SCATTER_REVERSE' sforward = 'SCATTER_FORWARD' diff --git a/devito/petsc/solve.py b/devito/petsc/solve.py index 1a3e2bbf95..c34b829798 100644 --- a/devito/petsc/solve.py +++ b/devito/petsc/solve.py @@ -1,17 +1,16 @@ from functools import singledispatch import sympy -from collections import defaultdict from devito.finite_differences.differentiable import Mul from devito.finite_differences.derivative import Derivative from devito.types import Eq, Symbol, SteppingDimension, TimeFunction from devito.types.equation import PetscEq from devito.operations.solve import eval_time_derivatives -from devito.symbolics import retrieve_functions, FieldFromComposite +from devito.symbolics import retrieve_functions from devito.tools import as_tuple, filter_ordered from devito.petsc.types import (LinearSolveExpr, PETScArray, DMDALocalInfo, - FieldData, MultipleFieldData, SubMatrices, PetscBundle) + FieldData, MultipleFieldData, SubMatrices) __all__ = ['PETScSolve', 'EssentialBC'] @@ -74,14 +73,14 @@ def generate_field_data(self, eqns, target, arrays): ) def build_function_eqns(self, eq, target, arrays): - b, F_target, zeroed_eqn, targets = separate_eqn(eq, target) + b, F_target, _, targets = separate_eqn(eq, target) formfunc = self.make_formfunc(eq, F_target, arrays, targets) formrhs = self.make_rhs(eq, b, arrays) return (formfunc, formrhs) def build_matvec_eqns(self, eq, target, arrays): - b, F_target, zeroed_eqn, targets = separate_eqn(eq, target) + b, F_target, _, targets = separate_eqn(eq, target) if not F_target: return None matvec = self.make_matvec(eq, F_target, arrays, targets) @@ -97,7 +96,9 @@ def make_matvec(self, eq, F_target, arrays, targets): def make_formfunc(self, eq, F_target, arrays, targets): if isinstance(eq, EssentialBC): - rhs = 0. + # TODO: CHECK THIS + rhs = arrays[main_target]['x'] - eq.rhs + # rhs = 0. else: if isinstance(F_target, (int, float)): rhs = F_target @@ -149,42 +150,19 @@ def linear_solve_args(self): arrays = self.generate_arrays_combined(*coupled_targets) - - # from IPython import embed; embed() - # TODO: don't need a 'b' for coupled - # arrays = self.generate_arrays_combined(*coupled_targets) - # all_formfuncs = [] - - # from IPython import embed; embed() - # f_field = [arrays[coupled_targets[t]]['f'] for t in range(len(coupled_targets))] - # f_field = PetscBundle(name='fuv_field', components=f_field) - - - all_data = MultipleFieldData(submatrices=jacobian, arrays=arrays, targets=coupled_targets) - - # f_v = arrays[coupled_targets[0]]['f'] - # f_u = arrays[coupled_targets[1]]['f'] - - + all_data = MultipleFieldData(submatrices=jacobian, arrays=arrays, + targets=coupled_targets) for target, eqns in self.target_eqns.items(): eqns = as_tuple(eqns) - - - # TODO: obvs fix and don't duplicate arrays here - # tmp_arr = self.generate_arrays(target) self.update_jacobian(eqns, target, jacobian, arrays[target]) - formfuncs = [self.build_function_eqns(eq, target, coupled_targets, arrays) for eq in eqns] - + formfuncs = [ + self.build_function_eqns(eq, target, coupled_targets, arrays) + for eq in eqns + ] all_data.extend_formfuncs(formfuncs) - # all_formfuncs.extend(formfuncs) - - # all_data. - - # from IPython import embed; embed() - return target, tuple(funcs), all_data def update_jacobian(self, eqns, target, jacobian, arrays): @@ -197,22 +175,10 @@ def update_jacobian(self, eqns, target, jacobian, arrays): if matvecs: jacobian.set_submatrix(target, submat, matvecs) - # def generate_field_data(self, eqns, target, arrays): - # # from IPython import embed; embed() - # formfuncs, formrhs = zip( - # *[self.build_function_eqns(eq, target, arrays) for eq in eqns] - # ) - - # return FieldData( - # target=target, - # formfuncs=formfuncs, - # formrhs=formrhs, - # arrays=arrays - # ) - def build_function_eqns(self, eq, main_target, coupled_targets, arrays): zeroed = eq.lhs - eq.rhs + # TODO: clean up, test coupled with time dependence zeroed_eqn = Eq(zeroed, 0) zeroed_eqn = eval_time_derivatives(zeroed) @@ -221,10 +187,9 @@ def build_function_eqns(self, eq, main_target, coupled_targets, arrays): target_funcs = generate_targets(Eq(zeroed, 0), t) mapper.update(targets_to_arrays(arrays[t]['x'], target_funcs)) - # from IPython import embed; embed() - if isinstance(eq, EssentialBC): - rhs = 0. + # TODO: CHECK THIS + rhs = arrays[main_target]['x'] - eq.rhs else: if isinstance(zeroed, (int, float)): rhs = zeroed @@ -233,7 +198,6 @@ def build_function_eqns(self, eq, main_target, coupled_targets, arrays): rhs = rhs.subs(self.time_mapper) return Eq(arrays[main_target]['f'], rhs, subdomain=eq.subdomain) - # return Eq(f_field, rhs, subdomain=eq.subdomain) def generate_arrays_combined(self, *targets): return { @@ -250,20 +214,6 @@ def generate_arrays_combined(self, *targets): } - # def make_formfunc(self, eq, F_target, arrays, targets): - # if isinstance(eq, EssentialBC): - # rhs = 0. - # else: - # if isinstance(F_target, (int, float)): - # rhs = F_target - # else: - # # from IPython import embed; embed() - # rhs = F_target.subs(targets_to_arrays(arrays['x'], targets)) - # rhs = rhs.subs(self.time_mapper) - # return Eq(arrays['f'], rhs, subdomain=eq.subdomain) - - - class EssentialBC(Eq): pass diff --git a/devito/petsc/types/array.py b/devito/petsc/types/array.py index 39eb2e2347..58ca3e28d7 100644 --- a/devito/petsc/types/array.py +++ b/devito/petsc/types/array.py @@ -1,15 +1,12 @@ from functools import cached_property from ctypes import POINTER, Structure -from sympy import Expr - from devito.types.utils import DimensionTuple -from devito.types.array import ArrayBasic, Bundle, ArrayMapped, ComponentAccess +from devito.types.array import ArrayBasic, Bundle, ComponentAccess from devito.finite_differences import Differentiable from devito.types.basic import AbstractFunction, IndexedData -from devito.tools import dtype_to_ctype, as_tuple, dtypes_vector_mapper, CustomDtype +from devito.tools import dtype_to_ctype, as_tuple from devito.symbolics import FieldFromComposite -from devito.petsc.types.object import PETScStruct class PETScArray(ArrayBasic, Differentiable): @@ -125,12 +122,11 @@ class PetscBundle(Bundle): """ """ is_Bundle = True - _data_alignment = False @property def _C_ctype(self): - # TODO: extend to cases with multiple petsc solves... + # TODO: extend to cases with multiple petsc solves...(need diff struct name for each solve) fields = [(i.target.name, dtype_to_ctype(i.dtype)) for i in self.components] return POINTER(type('Field', (Structure,), {'_fields_': fields})) @@ -142,11 +138,11 @@ def indexed(self): @cached_property def vector(self): return PETScArray( - name=self.name, - target=self.c0.target, - liveness=self.c0.liveness, - localinfo=self.c0.localinfo, - ) + name=self.name, + target=self.c0.target, + liveness=self.c0.liveness, + localinfo=self.c0.localinfo, + ) @property def _C_name(self): diff --git a/devito/petsc/types/object.py b/devito/petsc/types/object.py index 62bccc5b2c..7821e316d9 100644 --- a/devito/petsc/types/object.py +++ b/devito/petsc/types/object.py @@ -7,6 +7,12 @@ from devito.petsc.iet.utils import petsc_call +class PetscMixin: + @property + def _C_free_priority(self): + return FREE_PRIORITY[self] + + class CallbackDM(LocalObject): """ PETSc Data Management object (DM). This is the DM instance @@ -16,7 +22,7 @@ class CallbackDM(LocalObject): dtype = CustomDtype('DM') -class DM(LocalObject): +class DM(LocalObject, PetscMixin): """ PETSc Data Management object (DM). This is the primary DM instance created within the main kernel and linked to the SNES @@ -36,11 +42,6 @@ def dofs(self): def _C_free(self): return petsc_call('DMDestroy', [Byref(self.function)]) - # TODO: Switch to an enumeration? - @property - def _C_free_priority(self): - return 4 - DMCast = cast('DM') @@ -310,8 +311,6 @@ def _C_ctype(self): return POINTER(POINTER(c_char)) - -# class PetscArrayofStruct(LocalCompositeObject): -# pass -# """ -# """ \ No newline at end of file +FREE_PRIORITY = { + DM: 4, +} diff --git a/devito/petsc/types/types.py b/devito/petsc/types/types.py index e470986612..973d558f38 100644 --- a/devito/petsc/types/types.py +++ b/devito/petsc/types/types.py @@ -3,8 +3,6 @@ from devito.tools import Reconstructable, sympy_mutex, as_tuple from devito.tools.dtypes_lowering import dtype_mapper from devito.petsc.utils import petsc_variables -from devito.petsc.types.object import PETScStruct -# from devito.petsc.types.array import PetscArrayofStruct class MetaData(sympy.Function, Reconstructable): @@ -195,37 +193,17 @@ class MultipleFieldData(FieldData): def __init__(self, targets, arrays, submatrices=None): self._targets = as_tuple(targets) self._arrays = arrays - # self.field_data_list = [] self._submatrices = submatrices self._formfuncs = [] - # self._f_field = f_field - # def add_field_data(self, field_data): - # self.field_data_list.append(field_data) - - # def get_field_data(self, target): - # for field_data in self.field_data_list: - # if field_data.target == target: - # return field_data - # raise ValueError(f"FieldData with target {target} not found.") - pass + # pass def extend_formfuncs(self, formfuncs): self._formfuncs.extend(formfuncs) - # @property - # def target(self): - # return None - - # @property - # def targets(self): - # return tuple(field_data.target for field_data in self.field_data_list) - @property def space_dimensions(self): - # space_dims = {field_data.space_dimensions for field_data in self.field_data_list} space_dims = {t.space_dimensions for t in self.targets} - # from IPython import embed; embed() if len(space_dims) > 1: # TODO: This may not actually have to be the case, but enforcing it for now raise ValueError( @@ -260,10 +238,6 @@ def space_order(self): def submatrices(self): return self._submatrices - # @property - # def formfuncs(self): - # return self._formfuncs - @property def targets(self): return self._targets diff --git a/devito/tools/utils.py b/devito/tools/utils.py index 0a28de16a8..0657146040 100644 --- a/devito/tools/utils.py +++ b/devito/tools/utils.py @@ -346,3 +346,4 @@ def key(i): return (v, str(type(i))) return sorted(items, key=key, reverse=True) + diff --git a/devito/types/array.py b/devito/types/array.py index 32c8ccc97b..3c3cedfc88 100644 --- a/devito/types/array.py +++ b/devito/types/array.py @@ -495,7 +495,6 @@ def _C_get_field(self, region, dim, side=None): return self.c0._C_get_field(region, dim, side=side) def __getitem__(self, index): - # from IPython import embed; embed() index = as_tuple(index) if len(index) == self.ndim: return super().__getitem__(index) diff --git a/devito/types/basic.py b/devito/types/basic.py index 9990c55836..a28ca2f486 100644 --- a/devito/types/basic.py +++ b/devito/types/basic.py @@ -87,7 +87,6 @@ def _C_typedata(self): if isinstance(_type, CustomDtype): return _type - # from IPython import embed; embed() while issubclass(_type, _Pointer): _type = _type._type_ diff --git a/examples/petsc/random/biharmonic/02_biharmonic.py b/examples/petsc/random/biharmonic/02_biharmonic.py new file mode 100644 index 0000000000..8c15f3d3a2 --- /dev/null +++ b/examples/petsc/random/biharmonic/02_biharmonic.py @@ -0,0 +1,148 @@ +# ref - https://github.com/bueler/p4pdes/blob/master/c/ch7/biharm.c + + +import os +import numpy as np + +from devito import (Grid, Function, Eq, Operator, switchconfig, + configuration, SubDomain) + +from devito.petsc import PETScSolve, EssentialBC +from devito.petsc.initialize import PetscInitialize +configuration['compiler'] = 'custom' +os.environ['CC'] = 'mpicc' + +PetscInitialize() + +# Subdomains to implement BCs +class SubTop(SubDomain): + name = 'subtop' + + def define(self, dimensions): + x, y = dimensions + return {x: x, y: ('right', 1)} + + +class SubBottom(SubDomain): + name = 'subbottom' + + def define(self, dimensions): + x, y = dimensions + return {x: x, y: ('left', 1)} + + +class SubLeft(SubDomain): + name = 'subleft' + + def define(self, dimensions): + x, y = dimensions + return {x: ('left', 1), y: y} + + +class SubRight(SubDomain): + name = 'subright' + + def define(self, dimensions): + x, y = dimensions + return {x: ('right', 1), y: y} + + +def c(x): + return x**3 * (1 - x)**3 + +def ddc(x): + return 6.0 * x * (1 - x) * (1 - 5.0 * x + 5.0 * x**2) + +def d4c(x): + return -72.0 * (1 - 5.0 * x + 5.0 * x**2) + +def u_exact_fcn(x, y): + return c(x) * c(y) + +def lap_u_exact_fcn(x, y): + return -ddc(x) * c(y) - c(x) * ddc(y) + +def f_fcn(x, y): + return d4c(x) * c(y) + 2.0 * ddc(x) * ddc(y) + c(x) * d4c(y) + + +sub1 = SubTop() +sub2 = SubBottom() +sub3 = SubLeft() +sub4 = SubRight() + +subdomains = (sub1, sub2, sub3, sub4) + +Lx = np.float64(1.) +Ly = np.float64(1.) + +# n_values = [33, 53, 73, 93, 113] +n_values = [33] +dx = np.array([Lx/(n-1) for n in n_values]) + +u_errors = [] +v_errors = [] + +for n in n_values: + grid = Grid( + shape=(n, n), extent=(Lx, Ly), subdomains=subdomains, dtype=np.float64 + ) + + u = Function(name='u', grid=grid, space_order=2) + v = Function(name='v', grid=grid, space_order=2) + f = Function(name='f', grid=grid, space_order=2) + + u_exact = Function(name='u_exact', grid=grid, space_order=2) + lap_u = Function(name='lap_u', grid=grid, space_order=2) + + eqn1 = Eq(-v.laplace, f, subdomain=grid.interior) + eqn2 = Eq(-u.laplace, v, subdomain=grid.interior) + + tmpx = np.linspace(0, Lx, n).astype(np.float64) + tmpy = np.linspace(0, Ly, n).astype(np.float64) + X, Y = np.meshgrid(tmpx, tmpy) + + f.data[:] = f_fcn(X, Y) + + # # Create boundary condition expressions using subdomains + # TODO: add initial guess callback for mixed systems + bc_u = [EssentialBC(u, 0., subdomain=sub1)] + bc_u += [EssentialBC(u, 0., subdomain=sub2)] + bc_u += [EssentialBC(u, 0., subdomain=sub3)] + bc_u += [EssentialBC(u, 0., subdomain=sub4)] + bc_v = [EssentialBC(v, 0., subdomain=sub1)] + bc_v += [EssentialBC(v, 0., subdomain=sub2)] + bc_v += [EssentialBC(v, 0., subdomain=sub3)] + bc_v += [EssentialBC(v, 0., subdomain=sub4)] + + # T (see ref) is nonsymmetric so need to set default KSP type to GMRES + petsc = PETScSolve({v: [eqn1]+bc_v, u: [eqn2]+bc_u}, solver_parameters={'ksp_rtol': 1e-10}) + + with switchconfig(language='petsc'): + op = Operator(petsc) + op.apply() + + u_exact.data[:] = u_exact_fcn(X, Y) + lap_u.data[:] = lap_u_exact_fcn(X, Y) + + # Compute infinity norm for u + u_diff = u_exact.data[:] - u.data[:] + u_error = np.linalg.norm(u_diff.ravel(), ord=np.inf) / np.linalg.norm(u_exact.data[:].ravel(), ord=np.inf) + u_errors.append(u_error) + + # Compute infinity norm for lap_u + v_diff = lap_u.data[:] - v.data[:] + v_error = np.linalg.norm(v_diff.ravel(), ord=np.inf) / np.linalg.norm(lap_u.data[:].ravel(), ord=np.inf) + v_errors.append(v_error) + +# u_slope, _ = np.polyfit(np.log(dx), np.log(u_errors), 1) +# v_slope, _ = np.polyfit(np.log(dx), np.log(v_errors), 1) + +# assert u_slope > 1.9 +# assert u_slope < 2.1 + +# assert v_slope > 1.9 +# assert v_slope < 2.1 +print(op.ccode) +print("u_errors:", u_errors) +print("v_errors:", v_errors) diff --git a/examples/petsc/random/biharmonic/biharmonic_matfree.c b/examples/petsc/random/biharmonic/biharmonic_matfree.c new file mode 100644 index 0000000000..53121473e7 --- /dev/null +++ b/examples/petsc/random/biharmonic/biharmonic_matfree.c @@ -0,0 +1,556 @@ + +// # ref - https://github.com/bueler/p4pdes/blob/master/c/ch7/biharm.c + +static char help[] = +"Solve the linear biharmonic equation in 2D. Equation is\n" +" Lap^2 u = f\n" +"where Lap = - grad^2 is the positive Laplacian, equivalently\n" +" u_xxxx + 2 u_xxyy + u_yyyy = f(x,y)\n" +"Domain is unit square S = (0,1)^2. Boundary conditions are homogeneous\n" +"simply-supported: u = 0, Lap u = 0. The equation is rewritten as a\n" +"2x2 block system with SPD Laplacian blocks on the diagonal:\n" +" | Lap | 0 | | v | | f | \n" +" |-----|-----| |---| = |---| \n" +" | -I | Lap | | u | | 0 | \n" +"Includes manufactured, polynomial exact solution. The discretization is\n" +"structured-grid (DMDA) finite differences. Includes analytical Jacobian.\n" +"Recommended preconditioning combines fieldsplit:\n" +" -pc_type fieldsplit -pc_fieldsplit_type multiplicative|additive \n" +"with multigrid as the preconditioner for the diagonal blocks:\n" +" -fieldsplit_v_pc_type mg|gamg -fieldsplit_u_pc_type mg|gamg\n" +"(GMG requires setting levels and Galerkin coarsening.) One can also do\n" +"monolithic multigrid (-pc_type mg|gamg).\n\n"; + +#include + +typedef struct { + PetscReal v, u; +} Field; + +typedef struct { + PetscReal (*f)(PetscReal x, PetscReal y); // right-hand side +} BiharmCtx; + +struct JacobianCtx +{ + DM * subdms; + IS * fields; + Mat * submats; +} ; + +struct SubMatrixCtx +{ + IS * rows; + IS * cols; +} ; + +static PetscReal c(PetscReal x) { + return x*x*x * (1.0-x)*(1.0-x)*(1.0-x); +} + +static PetscReal ddc(PetscReal x) { + return 6.0 * x * (1.0-x) * (1.0 - 5.0 * x + 5.0 * x*x); +} + +static PetscReal d4c(PetscReal x) { + return - 72.0 * (1.0 - 5.0 * x + 5.0 * x*x); +} + +static PetscReal u_exact_fcn(PetscReal x, PetscReal y) { + return c(x) * c(y); +} + +static PetscReal lap_u_exact_fcn(PetscReal x, PetscReal y) { + return - ddc(x) * c(y) - c(x) * ddc(y); // Lap u = - grad^2 u +} + +static PetscReal f_fcn(PetscReal x, PetscReal y) { + return d4c(x) * c(y) + 2.0 * ddc(x) * ddc(y) + c(x) * d4c(y); // Lap^2 u = grad^4 u +} + +extern PetscErrorCode FormExactWLocal(DMDALocalInfo*, Field**, BiharmCtx*); +extern PetscErrorCode FormFunction(SNES snes, Vec X, Vec F, void* dummy); +extern PetscErrorCode J00_MatMult(Mat J, Vec X, Vec Y); +extern PetscErrorCode J10_MatMult(Mat J, Vec X, Vec Y); +extern PetscErrorCode J11_MatMult(Mat J, Vec X, Vec Y); +extern PetscErrorCode WholeMatMult(Mat J, Vec X, Vec Y); +PetscErrorCode MatCreateSubMatrices0(Mat J, PetscInt nfields, IS * irow, IS * icol, MatReuse scall, Mat * * submats); +extern PetscErrorCode PopulateMatContext(struct JacobianCtx * jctx, DM * subdms, IS * fields); + +int main(int argc,char **argv) { + DM da; + SNES snes; + Vec w, w_initial, w_exact; + BiharmCtx user; + Field **aW; + PetscReal normv, normu, errv, erru; + DMDALocalInfo info; + IS *fields; + DM *subdms; + PetscInt nfields; + + struct JacobianCtx jctx0; + Mat J; + + PetscCall(PetscInitialize(&argc,&argv,NULL,help)); + + user.f = &f_fcn; + PetscCall(DMDACreate2d(PETSC_COMM_WORLD, + DM_BOUNDARY_NONE, DM_BOUNDARY_NONE, DMDA_STENCIL_STAR, + 33,33,PETSC_DECIDE,PETSC_DECIDE, + 2,1, // degrees of freedom, stencil width + NULL,NULL,&da)); + PetscCall(DMSetApplicationContext(da,&user)); + PetscCall(DMSetFromOptions(da)); + PetscCall(DMSetUp(da)); // this must be called BEFORE SetUniformCoordinates + PetscCall(DMSetMatType(da, MATSHELL)); + PetscCall(DMDASetUniformCoordinates(da,0.0,1.0,0.0,1.0,-1.0,-1.0)); + PetscCall(DMDASetFieldName(da,0,"v")); + PetscCall(DMDASetFieldName(da,1,"u")); + PetscCall(DMCreateMatrix(da,&J)); + + PetscCall(SNESCreate(PETSC_COMM_WORLD,&snes)); + PetscCall(SNESSetDM(snes,da)); + PetscCall(SNESSetFunction(snes,NULL,FormFunction,NULL)); + PetscCall(SNESSetType(snes,SNESKSPONLY)); + PetscCall(SNESSetFromOptions(snes)); + + PetscCall(SNESSetJacobian(snes,J,J,MatMFFDComputeJacobian,NULL)); + PetscCall(MatShellSetOperation(J,MATOP_MULT,(void (*)(void))WholeMatMult)); + + PetscCall(MatSetDM(J,da)); + PetscCall(DMCreateFieldDecomposition(da,&(nfields),NULL,&fields,&subdms)); + PetscCall(PopulateMatContext(&(jctx0),subdms,fields)); + PetscCall(MatShellSetContext(J,&(jctx0))); + PetscCall(MatCreateSubMatrices0(J,nfields,fields,fields,MAT_INITIAL_MATRIX,&(jctx0.submats))); + + PetscCall(DMGetGlobalVector(da,&w_initial)); + PetscCall(VecSet(w_initial,0.0)); + PetscCall(SNESSolve(snes,NULL,w_initial)); + // PetscCall(VecView(w_initial,PETSC_VIEWER_STDOUT_WORLD)); + PetscCall(DMRestoreGlobalVector(da,&w_initial)); + PetscCall(DMDestroy(&da)); + + PetscCall(SNESGetSolution(snes,&w)); + PetscCall(SNESGetDM(snes,&da)); + PetscCall(DMDAGetLocalInfo(da,&info)); + + PetscCall(DMCreateGlobalVector(da,&w_exact)); + PetscCall(DMDAVecGetArray(da,w_exact,&aW)); + PetscCall(FormExactWLocal(&info,aW,&user)); + PetscCall(DMDAVecRestoreArray(da,w_exact,&aW)); + PetscCall(VecStrideNorm(w_exact,0,NORM_INFINITY,&normv)); + PetscCall(VecStrideNorm(w_exact,1,NORM_INFINITY,&normu)); + PetscCall(VecAXPY(w,-1.0,w_exact)); + PetscCall(VecStrideNorm(w,0,NORM_INFINITY,&errv)); + PetscCall(VecStrideNorm(w,1,NORM_INFINITY,&erru)); + PetscCall(PetscPrintf(PETSC_COMM_WORLD, + "done on %d x %d grid ...\n" + " errors |v-vex|_inf/|vex|_inf = %.5e, |u-uex|_inf/|uex|_inf = %.5e\n", + info.mx,info.my,errv/normv,erru/normu)); + + + PetscCall(ISDestroy(&(fields[0]))); + PetscCall(ISDestroy(&(fields[1]))); + PetscCall(PetscFree(fields)); + PetscCall(DMDestroy(&(subdms[0]))); + PetscCall(DMDestroy(&(subdms[1]))); + PetscCall(PetscFree(subdms)); + PetscCall(VecDestroy(&w_exact)); + PetscCall(MatDestroy(&J)); + PetscCall(SNESDestroy(&snes)); + PetscCall(PetscFinalize()); + return 0; +} + +PetscErrorCode FormExactWLocal(DMDALocalInfo *info, Field **aW, BiharmCtx *user) { + PetscInt i, j; + PetscReal xymin[2], xymax[2], hx, hy, x, y; + PetscCall(DMGetBoundingBox(info->da,xymin,xymax)); + hx = (xymax[0] - xymin[0]) / (info->mx - 1); + hy = (xymax[1] - xymin[1]) / (info->my - 1); + for (j = info->ys; j < info->ys + info->ym; j++) { + y = j * hy; + for (i = info->xs; i < info->xs + info->xm; i++) { + x = i * hx; + aW[j][i].u = u_exact_fcn(x,y); + aW[j][i].v = lap_u_exact_fcn(x,y); + } + } + return 0; +} + + +PetscErrorCode FormFunction(SNES snes, Vec X, Vec F, void * dummy) +{ + Vec xlocal, flocal; + DMDALocalInfo info; + DM da; + PetscScalar *x_vec, *f_vec; + + BiharmCtx *user; + + PetscCall(SNESGetDM(snes,&da)); + + PetscCall(DMGetApplicationContext(da,&user)); + + PetscCall(DMDAGetLocalInfo(da,&info)); + PetscCall(DMGetLocalVector(da,&xlocal)); + PetscCall(DMGetLocalVector(da,&flocal)); + + PetscCall(DMGlobalToLocalBegin(da,X,INSERT_VALUES,xlocal)); + PetscCall(DMGlobalToLocalEnd(da,X,INSERT_VALUES,xlocal)); + + PetscCall(VecGetArray(xlocal,&x_vec)); + PetscCall(VecGetArray(flocal,&f_vec)); + + Field (*xx)[info.gxm] = (Field (*)[info.gxm]) x_vec; + Field (*ff)[info.gxm] = (Field (*)[info.gxm]) f_vec; + + PetscInt i, j; + PetscReal xymin[2], xymax[2], hx, hy, darea, scx, scy, scdiag, x, y, + ve, vw, vn, vs, ue, uw, un, us; + + hx = 1. / (info.mx - 1); + hy = 1. / (info.my - 1); + + darea = hx * hy; // multiply FD equations by this + + scx = hy / hx; + scy = hx / hy; + scdiag = 2.0 * (scx + scy); // diagonal scaling + for (j = info.ys; j < info.ys + info.ym; j++) { + y = xymin[1] + j * hy; + for (i = info.xs; i < info.xs + info.xm; i++) { + x = xymin[0] + i * hx; + if (i==0 || i==info.mx-1 || j==0 || j==info.my-1) { + ff[j][i].v = scdiag * xx[j][i].v; + ff[j][i].u = scdiag * xx[j][i].u; + } else { + ve = xx[j][i+1].v; + vw = xx[j][i-1].v; + vn = xx[j+1][i].v; + vs = xx[j-1][i].v; + ff[j][i].v = scdiag * xx[j][i].v - scx * (vw + ve) - scy * (vs + vn) + - darea * (*(user->f))(x,y); + ue = xx[j][i+1].u; + uw = xx[j][i-1].u; + un = xx[j+1][i].u; + us = xx[j-1][i].u; + ff[j][i].u = - darea * xx[j][i].v + + scdiag * xx[j][i].u - scx * (uw + ue) - scy * (us + un); + } + } + } + + PetscCall(VecRestoreArray(xlocal,&x_vec)); + PetscCall(VecRestoreArray(flocal,&f_vec)); + + PetscCall(DMLocalToGlobalBegin(da,flocal,INSERT_VALUES,F)); + PetscCall(DMLocalToGlobalEnd(da,flocal,INSERT_VALUES,F)); + PetscCall(DMRestoreLocalVector(da,&xlocal)); + PetscCall(DMRestoreLocalVector(da,&flocal)); + + return 0; +} + + +PetscErrorCode J00_MatMult(Mat J, Vec X, Vec Y) +{ + PetscFunctionBeginUser; + + DM dm0; + DMDALocalInfo info; + Vec xloc; + Vec yloc; + + BiharmCtx * ctx0; + PetscScalar * x_v_vec; + PetscScalar * y_v_vec; + + PetscCall(MatGetDM(J,&(dm0))); + PetscCall(DMGetApplicationContext(dm0,&(ctx0))); + PetscCall(DMGetLocalVector(dm0,&(xloc))); + PetscCall(DMGlobalToLocalBegin(dm0,X,INSERT_VALUES,xloc)); + PetscCall(DMGlobalToLocalEnd(dm0,X,INSERT_VALUES,xloc)); + PetscCall(DMGetLocalVector(dm0,&(yloc))); + PetscCall(VecSet(yloc,0.0)); + PetscCall(VecGetArray(yloc,&y_v_vec)); + PetscCall(VecGetArray(xloc,&x_v_vec)); + PetscCall(DMDAGetLocalInfo(dm0,&(info))); + + PetscScalar (* x_v)[info.gxm] = (PetscScalar (*)[info.gxm]) x_v_vec; + PetscScalar (* y_v)[info.gxm] = (PetscScalar (*)[info.gxm]) y_v_vec; + + PetscInt i, j; + PetscReal xymin[2], xymax[2], hx, hy, darea, scx, scy, scdiag, x, y, + ve, vw, vn, vs, ue, uw, un, us; + + hx = 1./ (info.mx - 1); + hy = 1./ (info.my - 1); + darea = hx * hy; // multiply FD equations by this + scx = hy / hx; + scy = hx / hy; + scdiag = 2.0 * (scx + scy); // diagonal scaling + for (j = info.ys; j < info.ys + info.ym; j++) { + y = xymin[1] + j * hy; + for (i = info.xs; i < info.xs + info.xm; i++) { + x = xymin[0] + i * hx; + if (i==0 || i==info.mx-1 || j==0 || j==info.my-1) { + y_v[j][i] = scdiag * x_v[j][i]; + } else { + ve = x_v[j][i+1]; + vw = x_v[j][i-1]; + vn = x_v[j+1][i]; + vs = x_v[j-1][i]; + y_v[j][i] = scdiag * x_v[j][i] - scx * (vw + ve) - scy * (vs + vn); + + } + } + } + + PetscCall(VecRestoreArray(yloc,&y_v_vec)); + PetscCall(VecRestoreArray(xloc,&x_v_vec)); + PetscCall(DMLocalToGlobalBegin(dm0,yloc,ADD_VALUES,Y)); + PetscCall(DMLocalToGlobalEnd(dm0,yloc,ADD_VALUES,Y)); + PetscCall(DMRestoreLocalVector(dm0,&(xloc))); + PetscCall(DMRestoreLocalVector(dm0,&(yloc))); + + PetscFunctionReturn(0); +} + +PetscErrorCode J10_MatMult(Mat J, Vec X, Vec Y) +{ + PetscFunctionBeginUser; + + DM dm0; + DMDALocalInfo info; + Vec xloc; + Vec yloc; + + BiharmCtx * ctx0; + PetscScalar * x_v_vec; + PetscScalar * y_v_vec; + + PetscCall(MatGetDM(J,&(dm0))); + PetscCall(DMGetApplicationContext(dm0,&(ctx0))); + PetscCall(DMGetLocalVector(dm0,&(xloc))); + PetscCall(DMGlobalToLocalBegin(dm0,X,INSERT_VALUES,xloc)); + PetscCall(DMGlobalToLocalEnd(dm0,X,INSERT_VALUES,xloc)); + PetscCall(DMGetLocalVector(dm0,&(yloc))); + PetscCall(VecSet(yloc,0.0)); + PetscCall(VecGetArray(yloc,&y_v_vec)); + PetscCall(VecGetArray(xloc,&x_v_vec)); + PetscCall(DMDAGetLocalInfo(dm0,&(info))); + + PetscScalar (* x_v)[info.gxm] = (PetscScalar (*)[info.gxm]) x_v_vec; + PetscScalar (* y_v)[info.gxm] = (PetscScalar (*)[info.gxm]) y_v_vec; + + PetscInt i, j; + PetscReal xymin[2], xymax[2], hx, hy, darea, scx, scy, scdiag, x, y, + ve, vw, vn, vs, ue, uw, un, us; + + hx = 1. / (info.mx - 1); + hy = 1. / (info.my - 1); + + darea = hx * hy; // multiply FD equations by this + scx = hy / hx; + scy = hx / hy; + scdiag = 2.0 * (scx + scy); // diagonal scaling + for (j = info.ys; j < info.ys + info.ym; j++) { + y = xymin[1] + j * hy; + for (i = info.xs; i < info.xs + info.xm; i++) { + x = xymin[0] + i * hx; + if (i==0 || i==info.mx-1 || j==0 || j==info.my-1) { + y_v[j][i] = 0.0; + } else { + y_v[j][i] = -darea * x_v[j][i]; + + } + } + } + + PetscCall(VecRestoreArray(yloc,&y_v_vec)); + PetscCall(VecRestoreArray(xloc,&x_v_vec)); + PetscCall(DMLocalToGlobalBegin(dm0,yloc,ADD_VALUES,Y)); + PetscCall(DMLocalToGlobalEnd(dm0,yloc,ADD_VALUES,Y)); + PetscCall(DMRestoreLocalVector(dm0,&(xloc))); + PetscCall(DMRestoreLocalVector(dm0,&(yloc))); + + PetscFunctionReturn(0); +} + + +PetscErrorCode J11_MatMult(Mat J, Vec X, Vec Y) +{ + PetscFunctionBeginUser; + + DM dm0; + DMDALocalInfo info; + Vec xloc; + Vec yloc; + + BiharmCtx * ctx0; + PetscScalar * x_v_vec; + PetscScalar * y_v_vec; + + PetscCall(MatGetDM(J,&(dm0))); + PetscCall(DMGetApplicationContext(dm0,&(ctx0))); + PetscCall(DMGetLocalVector(dm0,&(xloc))); + PetscCall(DMGlobalToLocalBegin(dm0,X,INSERT_VALUES,xloc)); + PetscCall(DMGlobalToLocalEnd(dm0,X,INSERT_VALUES,xloc)); + PetscCall(DMGetLocalVector(dm0,&(yloc))); + PetscCall(VecSet(yloc,0.0)); + PetscCall(VecGetArray(yloc,&y_v_vec)); + PetscCall(VecGetArray(xloc,&x_v_vec)); + PetscCall(DMDAGetLocalInfo(dm0,&(info))); + + PetscScalar (* x_v)[info.gxm] = (PetscScalar (*)[info.gxm]) x_v_vec; + PetscScalar (* y_v)[info.gxm] = (PetscScalar (*)[info.gxm]) y_v_vec; + + PetscInt i, j; + PetscReal xymin[2], xymax[2], hx, hy, darea, scx, scy, scdiag, x, y, + ve, vw, vn, vs, ue, uw, un, us; +// PetscCall(DMGetBoundingBox(info.da,xymin,xymax)); + hx = 1. / (info.mx - 1); + hy = 1. / (info.my - 1); + + darea = hx * hy; + scx = hy / hx; + scy = hx / hy; + scdiag = 2.0 * (scx + scy); + for (j = info.ys; j < info.ys + info.ym; j++) { + y = xymin[1] + j * hy; + for (i = info.xs; i < info.xs + info.xm; i++) { + x = xymin[0] + i * hx; + if (i==0 || i==info.mx-1 || j==0 || j==info.my-1) { + y_v[j][i] = scdiag * x_v[j][i]; + } else { + ve = x_v[j][i+1]; + vw = x_v[j][i-1]; + vn = x_v[j+1][i]; + vs = x_v[j-1][i]; + y_v[j][i] = scdiag * x_v[j][i] - scx * (vw + ve) - scy * (vs + vn); + + } + } + } + + PetscCall(VecRestoreArray(yloc,&y_v_vec)); + PetscCall(VecRestoreArray(xloc,&x_v_vec)); + PetscCall(DMLocalToGlobalBegin(dm0,yloc,ADD_VALUES,Y)); + PetscCall(DMLocalToGlobalEnd(dm0,yloc,ADD_VALUES,Y)); + PetscCall(DMRestoreLocalVector(dm0,&(xloc))); + PetscCall(DMRestoreLocalVector(dm0,&(yloc))); + + PetscFunctionReturn(0); +} + +PetscErrorCode WholeMatMult(Mat J, Vec X, Vec Y) +{ + Vec J00X; + Vec J00Y; + Vec J10X; + Vec J10Y; + Vec J11X; + Vec J11Y; + + struct SubMatrixCtx * J00ctx; + struct SubMatrixCtx * J10ctx; + struct SubMatrixCtx * J11ctx; + struct JacobianCtx * jctx; + + PetscFunctionBeginUser; + + PetscCall(MatShellGetContext(J,&(jctx))); + + PetscCall(VecSet(Y,0.0)); + + Mat J00 = jctx->submats[0]; + PetscCall(MatShellGetContext(J00,&(J00ctx))); + PetscCall(VecGetSubVector(X,*(J00ctx->cols),&(J00X))); + PetscCall(VecGetSubVector(Y,*(J00ctx->rows),&(J00Y))); + PetscCall(MatMult(J00,J00X,J00Y)); + PetscCall(VecRestoreSubVector(X,*(J00ctx->cols),&(J00X))); + PetscCall(VecRestoreSubVector(Y,*(J00ctx->rows),&(J00Y))); + + Mat J10 = jctx->submats[2]; + PetscCall(MatShellGetContext(J10,&(J10ctx))); + PetscCall(VecGetSubVector(X,*(J10ctx->cols),&(J10X))); + PetscCall(VecGetSubVector(Y,*(J10ctx->rows),&(J10Y))); + PetscCall(MatMult(J10,J10X,J10Y)); + PetscCall(VecRestoreSubVector(X,*(J10ctx->cols),&(J10X))); + PetscCall(VecRestoreSubVector(Y,*(J10ctx->rows),&(J10Y))); + + Mat J11 = jctx->submats[3]; + PetscCall(MatShellGetContext(J11,&(J11ctx))); + PetscCall(VecGetSubVector(X,*(J11ctx->cols),&(J11X))); + PetscCall(VecGetSubVector(Y,*(J11ctx->rows),&(J11Y))); + PetscCall(MatMult(J11,J11X,J11Y)); + PetscCall(VecRestoreSubVector(X,*(J11ctx->cols),&(J11X))); + PetscCall(VecRestoreSubVector(Y,*(J11ctx->rows),&(J11Y))); + + + PetscFunctionReturn(0); +} + + +PetscErrorCode MatCreateSubMatrices0(Mat J, PetscInt nfields, IS * irow, IS * icol, MatReuse scall, Mat * * submats) +{ + PetscFunctionBeginUser; + + PetscInt M; + PetscInt N; + Mat block; + DM dm0; + PetscInt dof; + + struct UserCtx0 * ctx0; + struct JacobianCtx * jctx; + struct SubMatrixCtx * subctx; + + PetscCall(MatShellGetContext(J,&(jctx))); + DM * subdms = jctx->subdms; + + PetscInt nsubmats = nfields*nfields; + PetscCall(PetscCalloc1(nsubmats,submats)); + PetscCall(MatGetDM(J,&(dm0))); + PetscCall(DMGetApplicationContext(dm0,&(ctx0))); + PetscCall(DMDAGetInfo(dm0,NULL,&(M),&(N),NULL,NULL,NULL,NULL,&(dof),NULL,NULL,NULL,NULL,NULL)); + PetscInt subblockrows = M*N; + PetscInt subblockcols = M*N; + Mat * submat_arr = *submats; + + for (int i = 0; i <= nsubmats - 1; i += 1) + { + PetscCall(MatCreate(PETSC_COMM_WORLD,&(block))); + PetscCall(MatSetSizes(block,PETSC_DECIDE,PETSC_DECIDE,subblockrows,subblockcols)); + PetscCall(MatSetType(block,MATSHELL)); + PetscCall(PetscMalloc1(1,&(subctx))); + PetscInt rowidx = i / dof; + PetscInt colidx = (i)%(dof); + subctx->rows = &(irow[rowidx]); + subctx->cols = &(icol[colidx]); + PetscCall(DMSetApplicationContext(subdms[rowidx],ctx0)); + PetscCall(MatSetDM(block,subdms[rowidx])); + PetscCall(MatShellSetContext(block,subctx)); + PetscCall(MatSetUp(block)); + submat_arr[i] = block; + } + PetscCall(MatShellSetOperation(submat_arr[0],MATOP_MULT,(void (*)(void))J00_MatMult)); + PetscCall(MatShellSetOperation(submat_arr[2],MATOP_MULT,(void (*)(void))J10_MatMult)); + PetscCall(MatShellSetOperation(submat_arr[3],MATOP_MULT,(void (*)(void))J11_MatMult)); + + PetscFunctionReturn(0); +} + + +PetscErrorCode PopulateMatContext(struct JacobianCtx * jctx, DM * subdms, IS * fields) +{ + PetscFunctionBeginUser; + + jctx->subdms = subdms; + jctx->fields = fields; + + PetscFunctionReturn(0); +} diff --git a/examples/petsc/random/biharmonic/biharmonic_matfree_nonscaled.c b/examples/petsc/random/biharmonic/biharmonic_matfree_nonscaled.c new file mode 100644 index 0000000000..04f2a02927 --- /dev/null +++ b/examples/petsc/random/biharmonic/biharmonic_matfree_nonscaled.c @@ -0,0 +1,553 @@ + +// # ref - https://github.com/bueler/p4pdes/blob/master/c/ch7/biharm.c + +static char help[] = +"Solve the linear biharmonic equation in 2D. Equation is\n" +" Lap^2 u = f\n" +"where Lap = - grad^2 is the positive Laplacian, equivalently\n" +" u_xxxx + 2 u_xxyy + u_yyyy = f(x,y)\n" +"Domain is unit square S = (0,1)^2. Boundary conditions are homogeneous\n" +"simply-supported: u = 0, Lap u = 0. The equation is rewritten as a\n" +"2x2 block system with SPD Laplacian blocks on the diagonal:\n" +" | Lap | 0 | | v | | f | \n" +" |-----|-----| |---| = |---| \n" +" | -I | Lap | | u | | 0 | \n" +"Includes manufactured, polynomial exact solution. The discretization is\n" +"structured-grid (DMDA) finite differences. Includes analytical Jacobian.\n" +"Recommended preconditioning combines fieldsplit:\n" +" -pc_type fieldsplit -pc_fieldsplit_type multiplicative|additive \n" +"with multigrid as the preconditioner for the diagonal blocks:\n" +" -fieldsplit_v_pc_type mg|gamg -fieldsplit_u_pc_type mg|gamg\n" +"(GMG requires setting levels and Galerkin coarsening.) One can also do\n" +"monolithic multigrid (-pc_type mg|gamg).\n\n"; + +#include + +typedef struct { + PetscReal v, u; +} Field; + +typedef struct { + PetscReal (*f)(PetscReal x, PetscReal y); // right-hand side +} BiharmCtx; + +struct JacobianCtx +{ + DM * subdms; + IS * fields; + Mat * submats; +} ; + +struct SubMatrixCtx +{ + IS * rows; + IS * cols; +} ; + +static PetscReal c(PetscReal x) { + return x*x*x * (1.0-x)*(1.0-x)*(1.0-x); +} + +static PetscReal ddc(PetscReal x) { + return 6.0 * x * (1.0-x) * (1.0 - 5.0 * x + 5.0 * x*x); +} + +static PetscReal d4c(PetscReal x) { + return - 72.0 * (1.0 - 5.0 * x + 5.0 * x*x); +} + +static PetscReal u_exact_fcn(PetscReal x, PetscReal y) { + return c(x) * c(y); +} + +static PetscReal lap_u_exact_fcn(PetscReal x, PetscReal y) { + return - ddc(x) * c(y) - c(x) * ddc(y); // Lap u = - grad^2 u +} + +static PetscReal f_fcn(PetscReal x, PetscReal y) { + return d4c(x) * c(y) + 2.0 * ddc(x) * ddc(y) + c(x) * d4c(y); // Lap^2 u = grad^4 u +} + +extern PetscErrorCode FormExactWLocal(DMDALocalInfo*, Field**, BiharmCtx*); +extern PetscErrorCode FormFunction(SNES snes, Vec X, Vec F, void* dummy); +extern PetscErrorCode J00_MatMult(Mat J, Vec X, Vec Y); +extern PetscErrorCode J10_MatMult(Mat J, Vec X, Vec Y); +extern PetscErrorCode J11_MatMult(Mat J, Vec X, Vec Y); +extern PetscErrorCode WholeMatMult(Mat J, Vec X, Vec Y); +PetscErrorCode MatCreateSubMatrices0(Mat J, PetscInt nfields, IS * irow, IS * icol, MatReuse scall, Mat * * submats); +extern PetscErrorCode PopulateMatContext(struct JacobianCtx * jctx, DM * subdms, IS * fields); + +int main(int argc,char **argv) { + DM da; + SNES snes; + Vec w, w_initial, w_exact; + BiharmCtx user; + Field **aW; + PetscReal normv, normu, errv, erru; + DMDALocalInfo info; + IS *fields; + DM *subdms; + PetscInt nfields; + + struct JacobianCtx jctx0; + Mat J; + + PetscCall(PetscInitialize(&argc,&argv,NULL,help)); + + user.f = &f_fcn; + PetscCall(DMDACreate2d(PETSC_COMM_WORLD, + DM_BOUNDARY_NONE, DM_BOUNDARY_NONE, DMDA_STENCIL_STAR, + 33,33,PETSC_DECIDE,PETSC_DECIDE, + 2,1, // degrees of freedom, stencil width + NULL,NULL,&da)); + PetscCall(DMSetApplicationContext(da,&user)); + PetscCall(DMSetFromOptions(da)); + PetscCall(DMSetUp(da)); // this must be called BEFORE SetUniformCoordinates + PetscCall(DMSetMatType(da, MATSHELL)); + PetscCall(DMDASetUniformCoordinates(da,0.0,1.0,0.0,1.0,-1.0,-1.0)); + PetscCall(DMDASetFieldName(da,0,"v")); + PetscCall(DMDASetFieldName(da,1,"u")); + PetscCall(DMCreateMatrix(da,&J)); + + PetscCall(SNESCreate(PETSC_COMM_WORLD,&snes)); + PetscCall(SNESSetDM(snes,da)); + PetscCall(SNESSetFunction(snes,NULL,FormFunction,NULL)); + PetscCall(SNESSetType(snes,SNESKSPONLY)); + PetscCall(SNESSetFromOptions(snes)); + + PetscCall(SNESSetJacobian(snes,J,J,MatMFFDComputeJacobian,NULL)); + PetscCall(MatShellSetOperation(J,MATOP_MULT,(void (*)(void))WholeMatMult)); + + PetscCall(MatSetDM(J,da)); + PetscCall(DMCreateFieldDecomposition(da,&(nfields),NULL,&fields,&subdms)); + PetscCall(PopulateMatContext(&(jctx0),subdms,fields)); + PetscCall(MatShellSetContext(J,&(jctx0))); + PetscCall(MatCreateSubMatrices0(J,nfields,fields,fields,MAT_INITIAL_MATRIX,&(jctx0.submats))); + + PetscCall(DMGetGlobalVector(da,&w_initial)); + PetscCall(VecSet(w_initial,0.0)); + PetscCall(SNESSolve(snes,NULL,w_initial)); + // PetscCall(VecView(w_initial,PETSC_VIEWER_STDOUT_WORLD)); + PetscCall(DMRestoreGlobalVector(da,&w_initial)); + PetscCall(DMDestroy(&da)); + + PetscCall(SNESGetSolution(snes,&w)); + PetscCall(SNESGetDM(snes,&da)); + PetscCall(DMDAGetLocalInfo(da,&info)); + + PetscCall(DMCreateGlobalVector(da,&w_exact)); + PetscCall(DMDAVecGetArray(da,w_exact,&aW)); + PetscCall(FormExactWLocal(&info,aW,&user)); + + PetscCall(DMDAVecRestoreArray(da,w_exact,&aW)); + // PetscCall(VecView(w_exact,PETSC_VIEWER_STDOUT_WORLD)); + PetscCall(VecStrideNorm(w_exact,0,NORM_INFINITY,&normv)); + PetscCall(VecStrideNorm(w_exact,1,NORM_INFINITY,&normu)); + PetscCall(VecAXPY(w,-1.0,w_exact)); + PetscCall(VecStrideNorm(w,0,NORM_INFINITY,&errv)); + PetscCall(VecStrideNorm(w,1,NORM_INFINITY,&erru)); + PetscCall(PetscPrintf(PETSC_COMM_WORLD, + "done on %d x %d grid ...\n" + " errors |v-vex|_inf/|vex|_inf = %.5e, |u-uex|_inf/|uex|_inf = %.5e\n", + info.mx,info.my,errv/normv,erru/normu)); + + + PetscCall(ISDestroy(&(fields[0]))); + PetscCall(ISDestroy(&(fields[1]))); + PetscCall(PetscFree(fields)); + PetscCall(DMDestroy(&(subdms[0]))); + PetscCall(DMDestroy(&(subdms[1]))); + PetscCall(PetscFree(subdms)); + PetscCall(VecDestroy(&w_exact)); + PetscCall(MatDestroy(&J)); + PetscCall(SNESDestroy(&snes)); + PetscCall(PetscFinalize()); + return 0; +} + +PetscErrorCode FormExactWLocal(DMDALocalInfo *info, Field **aW, BiharmCtx *user) { + PetscInt i, j; + PetscReal xymin[2], xymax[2], hx, hy, x, y; + PetscCall(DMGetBoundingBox(info->da,xymin,xymax)); + hx = (xymax[0] - xymin[0]) / (info->mx - 1); + hy = (xymax[1] - xymin[1]) / (info->my - 1); + for (j = info->ys; j < info->ys + info->ym; j++) { + y = j * hy; + for (i = info->xs; i < info->xs + info->xm; i++) { + x = i * hx; + aW[j][i].u = u_exact_fcn(x,y); + aW[j][i].v = lap_u_exact_fcn(x,y); + } + } + return 0; +} + + +PetscErrorCode FormFunction(SNES snes, Vec X, Vec F, void * dummy) +{ + Vec xlocal, flocal; + DMDALocalInfo info; + DM da; + PetscScalar *x_vec, *f_vec; + + BiharmCtx *user; + + PetscCall(SNESGetDM(snes,&da)); + + PetscCall(DMGetApplicationContext(da,&user)); + + PetscCall(DMDAGetLocalInfo(da,&info)); + PetscCall(DMGetLocalVector(da,&xlocal)); + PetscCall(DMGetLocalVector(da,&flocal)); + + PetscCall(DMGlobalToLocalBegin(da,X,INSERT_VALUES,xlocal)); + PetscCall(DMGlobalToLocalEnd(da,X,INSERT_VALUES,xlocal)); + + PetscCall(VecGetArray(xlocal,&x_vec)); + PetscCall(VecGetArray(flocal,&f_vec)); + + Field (*xx)[info.gxm] = (Field (*)[info.gxm]) x_vec; + Field (*ff)[info.gxm] = (Field (*)[info.gxm]) f_vec; + + PetscInt i, j; + PetscReal xymin[2], xymax[2], hx, hy, darea, scx, scy, scdiag, x, y, + ve, vw, vn, vs, ue, uw, un, us; + + hx = 1. / (info.mx - 1); + hy = 1. / (info.my - 1); + + darea = hx * hy; // multiply FD equations by this + + scx = 1. / (hx*hx); + scy = 1. / (hy*hy); + scdiag = 2.0 * (scx + scy); // diagonal scaling + for (j = info.ys; j < info.ys + info.ym; j++) { + y = xymin[1] + j * hy; + for (i = info.xs; i < info.xs + info.xm; i++) { + x = xymin[0] + i * hx; + if (i==0 || i==info.mx-1 || j==0 || j==info.my-1) { + ff[j][i].v = xx[j][i].v; + ff[j][i].u = xx[j][i].u; + } else { + ve = xx[j][i+1].v; + vw = xx[j][i-1].v; + vn = xx[j+1][i].v; + vs = xx[j-1][i].v; + ff[j][i].v = scdiag * xx[j][i].v - scx * (vw + ve) - scy * (vs + vn) + - (*(user->f))(x,y); + ue = xx[j][i+1].u; + uw = xx[j][i-1].u; + un = xx[j+1][i].u; + us = xx[j-1][i].u; + ff[j][i].u = -xx[j][i].v + + scdiag * xx[j][i].u - scx * (uw + ue) - scy * (us + un); + } + } + } + + PetscCall(VecRestoreArray(xlocal,&x_vec)); + PetscCall(VecRestoreArray(flocal,&f_vec)); + + PetscCall(DMLocalToGlobalBegin(da,flocal,INSERT_VALUES,F)); + PetscCall(DMLocalToGlobalEnd(da,flocal,INSERT_VALUES,F)); + PetscCall(DMRestoreLocalVector(da,&xlocal)); + PetscCall(DMRestoreLocalVector(da,&flocal)); + + return 0; +} + + +PetscErrorCode J00_MatMult(Mat J, Vec X, Vec Y) +{ + PetscFunctionBeginUser; + + DM dm0; + DMDALocalInfo info; + Vec xloc; + Vec yloc; + + BiharmCtx * ctx0; + PetscScalar * x_v_vec; + PetscScalar * y_v_vec; + + PetscCall(MatGetDM(J,&(dm0))); + PetscCall(DMGetApplicationContext(dm0,&(ctx0))); + PetscCall(DMGetLocalVector(dm0,&(xloc))); + PetscCall(DMGlobalToLocalBegin(dm0,X,INSERT_VALUES,xloc)); + PetscCall(DMGlobalToLocalEnd(dm0,X,INSERT_VALUES,xloc)); + PetscCall(DMGetLocalVector(dm0,&(yloc))); + PetscCall(VecSet(yloc,0.0)); + PetscCall(VecGetArray(yloc,&y_v_vec)); + PetscCall(VecGetArray(xloc,&x_v_vec)); + PetscCall(DMDAGetLocalInfo(dm0,&(info))); + + PetscScalar (* x_v)[info.gxm] = (PetscScalar (*)[info.gxm]) x_v_vec; + PetscScalar (* y_v)[info.gxm] = (PetscScalar (*)[info.gxm]) y_v_vec; + + PetscInt i, j; + PetscReal xymin[2], xymax[2], hx, hy, darea, scx, scy, scdiag, x, y, + ve, vw, vn, vs, ue, uw, un, us; + + hx = 1./ (info.mx - 1); + hy = 1./ (info.my - 1); + darea = hx * hy; // multiply FD equations by this + scx = 1. / (hx*hx); + scy = 1. / (hy*hy); + scdiag = 2.0 * (scx + scy); // diagonal scaling + for (j = info.ys; j < info.ys + info.ym; j++) { + y = xymin[1] + j * hy; + for (i = info.xs; i < info.xs + info.xm; i++) { + x = xymin[0] + i * hx; + if (i==0 || i==info.mx-1 || j==0 || j==info.my-1) { + y_v[j][i] = x_v[j][i]; + } else { + ve = x_v[j][i+1]; + vw = x_v[j][i-1]; + vn = x_v[j+1][i]; + vs = x_v[j-1][i]; + y_v[j][i] = scdiag * x_v[j][i] - scx * (vw + ve) - scy * (vs + vn); + + } + } + } + + PetscCall(VecRestoreArray(yloc,&y_v_vec)); + PetscCall(VecRestoreArray(xloc,&x_v_vec)); + PetscCall(DMLocalToGlobalBegin(dm0,yloc,ADD_VALUES,Y)); + PetscCall(DMLocalToGlobalEnd(dm0,yloc,ADD_VALUES,Y)); + PetscCall(DMRestoreLocalVector(dm0,&(xloc))); + PetscCall(DMRestoreLocalVector(dm0,&(yloc))); + + PetscFunctionReturn(0); +} + +PetscErrorCode J10_MatMult(Mat J, Vec X, Vec Y) +{ + PetscFunctionBeginUser; + + DM dm0; + DMDALocalInfo info; + Vec xloc; + Vec yloc; + + BiharmCtx * ctx0; + PetscScalar * x_v_vec; + PetscScalar * y_v_vec; + + PetscCall(MatGetDM(J,&(dm0))); + PetscCall(DMGetApplicationContext(dm0,&(ctx0))); + PetscCall(DMGetLocalVector(dm0,&(xloc))); + PetscCall(DMGlobalToLocalBegin(dm0,X,INSERT_VALUES,xloc)); + PetscCall(DMGlobalToLocalEnd(dm0,X,INSERT_VALUES,xloc)); + PetscCall(DMGetLocalVector(dm0,&(yloc))); + PetscCall(VecSet(yloc,0.0)); + PetscCall(VecGetArray(yloc,&y_v_vec)); + PetscCall(VecGetArray(xloc,&x_v_vec)); + PetscCall(DMDAGetLocalInfo(dm0,&(info))); + + PetscScalar (* x_v)[info.gxm] = (PetscScalar (*)[info.gxm]) x_v_vec; + PetscScalar (* y_v)[info.gxm] = (PetscScalar (*)[info.gxm]) y_v_vec; + + PetscInt i, j; + PetscReal xymin[2], xymax[2], hx, hy, darea, scx, scy, scdiag, x, y, + ve, vw, vn, vs, ue, uw, un, us; + + hx = 1. / (info.mx - 1); + hy = 1. / (info.my - 1); + + darea = hx * hy; // multiply FD equations by this + scx = 1. / (hx*hx); + scy = 1. / (hy*hy); + scdiag = 2.0 * (scx + scy); // diagonal scaling + // print info.ys to screen + // PetscCall(PetscPrintf(PETSC_COMM_WORLD,"info.ys = %d\n",info.xm)); + for (j = info.ys+1; j < info.ys + info.ym-1; j++) { + for (i = info.xs+1; i < info.xs + info.xm-1; i++) { + y_v[j][i] = -x_v[j][i]; + } + } + + PetscCall(VecRestoreArray(yloc,&y_v_vec)); + PetscCall(VecRestoreArray(xloc,&x_v_vec)); + PetscCall(DMLocalToGlobalBegin(dm0,yloc,ADD_VALUES,Y)); + PetscCall(DMLocalToGlobalEnd(dm0,yloc,ADD_VALUES,Y)); + PetscCall(DMRestoreLocalVector(dm0,&(xloc))); + PetscCall(DMRestoreLocalVector(dm0,&(yloc))); + + PetscFunctionReturn(0); +} + + +PetscErrorCode J11_MatMult(Mat J, Vec X, Vec Y) +{ + PetscFunctionBeginUser; + + DM dm0; + DMDALocalInfo info; + Vec xloc; + Vec yloc; + + BiharmCtx * ctx0; + PetscScalar * x_v_vec; + PetscScalar * y_v_vec; + + PetscCall(MatGetDM(J,&(dm0))); + PetscCall(DMGetApplicationContext(dm0,&(ctx0))); + PetscCall(DMGetLocalVector(dm0,&(xloc))); + PetscCall(DMGlobalToLocalBegin(dm0,X,INSERT_VALUES,xloc)); + PetscCall(DMGlobalToLocalEnd(dm0,X,INSERT_VALUES,xloc)); + PetscCall(DMGetLocalVector(dm0,&(yloc))); + PetscCall(VecSet(yloc,0.0)); + PetscCall(VecGetArray(yloc,&y_v_vec)); + PetscCall(VecGetArray(xloc,&x_v_vec)); + PetscCall(DMDAGetLocalInfo(dm0,&(info))); + + PetscScalar (* x_v)[info.gxm] = (PetscScalar (*)[info.gxm]) x_v_vec; + PetscScalar (* y_v)[info.gxm] = (PetscScalar (*)[info.gxm]) y_v_vec; + + PetscInt i, j; + PetscReal xymin[2], xymax[2], hx, hy, darea, scx, scy, scdiag, x, y, + ve, vw, vn, vs, ue, uw, un, us; +// PetscCall(DMGetBoundingBox(info.da,xymin,xymax)); + hx = 1. / (info.mx - 1); + hy = 1. / (info.my - 1); + + darea = hx * hy; + scx = 1. / (hx*hx); + scy = 1. / (hy*hy); + scdiag = 2.0 * (scx + scy); + for (j = info.ys; j < info.ys + info.ym; j++) { + y = xymin[1] + j * hy; + for (i = info.xs; i < info.xs + info.xm; i++) { + x = xymin[0] + i * hx; + if (i==0 || i==info.mx-1 || j==0 || j==info.my-1) { + y_v[j][i] = x_v[j][i]; + } else { + ve = x_v[j][i+1]; + vw = x_v[j][i-1]; + vn = x_v[j+1][i]; + vs = x_v[j-1][i]; + y_v[j][i] = scdiag * x_v[j][i] - scx * (vw + ve) - scy * (vs + vn); + + } + } + } + + PetscCall(VecRestoreArray(yloc,&y_v_vec)); + PetscCall(VecRestoreArray(xloc,&x_v_vec)); + PetscCall(DMLocalToGlobalBegin(dm0,yloc,ADD_VALUES,Y)); + PetscCall(DMLocalToGlobalEnd(dm0,yloc,ADD_VALUES,Y)); + PetscCall(DMRestoreLocalVector(dm0,&(xloc))); + PetscCall(DMRestoreLocalVector(dm0,&(yloc))); + + PetscFunctionReturn(0); +} + +PetscErrorCode WholeMatMult(Mat J, Vec X, Vec Y) +{ + Vec J00X; + Vec J00Y; + Vec J10X; + Vec J10Y; + Vec J11X; + Vec J11Y; + + struct SubMatrixCtx * J00ctx; + struct SubMatrixCtx * J10ctx; + struct SubMatrixCtx * J11ctx; + struct JacobianCtx * jctx; + + PetscFunctionBeginUser; + + PetscCall(MatShellGetContext(J,&(jctx))); + + PetscCall(VecSet(Y,0.0)); + + Mat J00 = jctx->submats[0]; + PetscCall(MatShellGetContext(J00,&(J00ctx))); + PetscCall(VecGetSubVector(X,*(J00ctx->cols),&(J00X))); + PetscCall(VecGetSubVector(Y,*(J00ctx->rows),&(J00Y))); + PetscCall(MatMult(J00,J00X,J00Y)); + PetscCall(VecRestoreSubVector(X,*(J00ctx->cols),&(J00X))); + PetscCall(VecRestoreSubVector(Y,*(J00ctx->rows),&(J00Y))); + + Mat J10 = jctx->submats[2]; + PetscCall(MatShellGetContext(J10,&(J10ctx))); + PetscCall(VecGetSubVector(X,*(J10ctx->cols),&(J10X))); + PetscCall(VecGetSubVector(Y,*(J10ctx->rows),&(J10Y))); + PetscCall(MatMult(J10,J10X,J10Y)); + PetscCall(VecRestoreSubVector(X,*(J10ctx->cols),&(J10X))); + PetscCall(VecRestoreSubVector(Y,*(J10ctx->rows),&(J10Y))); + + Mat J11 = jctx->submats[3]; + PetscCall(MatShellGetContext(J11,&(J11ctx))); + PetscCall(VecGetSubVector(X,*(J11ctx->cols),&(J11X))); + PetscCall(VecGetSubVector(Y,*(J11ctx->rows),&(J11Y))); + PetscCall(MatMult(J11,J11X,J11Y)); + PetscCall(VecRestoreSubVector(X,*(J11ctx->cols),&(J11X))); + PetscCall(VecRestoreSubVector(Y,*(J11ctx->rows),&(J11Y))); + + + PetscFunctionReturn(0); +} + + +PetscErrorCode MatCreateSubMatrices0(Mat J, PetscInt nfields, IS * irow, IS * icol, MatReuse scall, Mat * * submats) +{ + PetscFunctionBeginUser; + + PetscInt M; + PetscInt N; + Mat block; + DM dm0; + PetscInt dof; + + struct UserCtx0 * ctx0; + struct JacobianCtx * jctx; + struct SubMatrixCtx * subctx; + + PetscCall(MatShellGetContext(J,&(jctx))); + DM * subdms = jctx->subdms; + + PetscInt nsubmats = nfields*nfields; + PetscCall(PetscCalloc1(nsubmats,submats)); + PetscCall(MatGetDM(J,&(dm0))); + PetscCall(DMGetApplicationContext(dm0,&(ctx0))); + PetscCall(DMDAGetInfo(dm0,NULL,&(M),&(N),NULL,NULL,NULL,NULL,&(dof),NULL,NULL,NULL,NULL,NULL)); + PetscInt subblockrows = M*N; + PetscInt subblockcols = M*N; + Mat * submat_arr = *submats; + + for (int i = 0; i <= nsubmats - 1; i += 1) + { + PetscCall(MatCreate(PETSC_COMM_WORLD,&(block))); + PetscCall(MatSetSizes(block,PETSC_DECIDE,PETSC_DECIDE,subblockrows,subblockcols)); + PetscCall(MatSetType(block,MATSHELL)); + PetscCall(PetscMalloc1(1,&(subctx))); + PetscInt rowidx = i / dof; + PetscInt colidx = (i)%(dof); + subctx->rows = &(irow[rowidx]); + subctx->cols = &(icol[colidx]); + PetscCall(DMSetApplicationContext(subdms[rowidx],ctx0)); + PetscCall(MatSetDM(block,subdms[rowidx])); + PetscCall(MatShellSetContext(block,subctx)); + PetscCall(MatSetUp(block)); + submat_arr[i] = block; + } + PetscCall(MatShellSetOperation(submat_arr[0],MATOP_MULT,(void (*)(void))J00_MatMult)); + PetscCall(MatShellSetOperation(submat_arr[2],MATOP_MULT,(void (*)(void))J10_MatMult)); + PetscCall(MatShellSetOperation(submat_arr[3],MATOP_MULT,(void (*)(void))J11_MatMult)); + + PetscFunctionReturn(0); +} + + +PetscErrorCode PopulateMatContext(struct JacobianCtx * jctx, DM * subdms, IS * fields) +{ + PetscFunctionBeginUser; + + jctx->subdms = subdms; + jctx->fields = fields; + + PetscFunctionReturn(0); +} diff --git a/tests/test_petsc.py b/tests/test_petsc.py index e89f4c4476..943331f3c3 100644 --- a/tests/test_petsc.py +++ b/tests/test_petsc.py @@ -1097,3 +1097,7 @@ def test_submatrices(self): # add tests for all new callbacks # def test_create_whole_matmult(): + + +# // add coupled test with 3 targets +# // add coupled test with 1 target? From 1ddc053f3d8565deaecd73dfa1936d36a9edfe43 Mon Sep 17 00:00:00 2001 From: ZoeLeibowitz Date: Thu, 29 May 2025 00:24:42 +0100 Subject: [PATCH 28/84] compiler: Working reuse_efunc for petscbundles --- devito/passes/iet/engine.py | 27 ++++++--------------------- devito/petsc/types/object.py | 3 --- devito/types/object.py | 2 ++ 3 files changed, 8 insertions(+), 24 deletions(-) diff --git a/devito/passes/iet/engine.py b/devito/passes/iet/engine.py index a90f4ba275..0c3c7481bf 100644 --- a/devito/passes/iet/engine.py +++ b/devito/passes/iet/engine.py @@ -151,7 +151,7 @@ def apply(self, func, **kwargs): if len(efuncs) > len(self.efuncs): efuncs = reuse_compounds(efuncs, self.sregistry) # TODO: fix for petsc bundles - # efuncs = reuse_efuncs(self.root, efuncs, self.sregistry) + efuncs = reuse_efuncs(self.root, efuncs, self.sregistry) self.efuncs = efuncs @@ -355,7 +355,7 @@ def reuse_efuncs(root, efuncs, sregistry=None): if isinstance(efunc, AsyncCallable): mapper[len(mapper)] = (efunc, [efunc]) continue - # from IPython import embed; embed() + afunc = abstract_efunc(efunc) key = afunc._signature() @@ -396,7 +396,7 @@ def abstract_efunc(efunc): functions = FindSymbols('basics|symbolics|dimensions').visit(efunc) mapper = abstract_objects(functions) - # from IPython import embed; embed() + efunc = Uxreplace(mapper).visit(efunc) efunc = efunc._rebuild(name='foo') @@ -416,8 +416,8 @@ def abstract_objects(objects0, sregistry=None): # Precedence rules make it possible to reconstruct objects that depend on # higher priority objects - # keys = [Bundle, Array, PETScArray, DiscreteFunction, AbstractIncrDimension, BlockDimension] - keys = [Bundle, PETScArray, DiscreteFunction, AbstractIncrDimension, BlockDimension] + keys = [Bundle, Array, PETScArray, DiscreteFunction, AbstractIncrDimension, BlockDimension] + # keys = [Bundle, Array, DiscreteFunction, AbstractIncrDimension, BlockDimension] priority = {k: i for i, k in enumerate(keys, start=1)} objects = sorted_priority(objects, priority) @@ -428,7 +428,6 @@ def abstract_objects(objects0, sregistry=None): for i in objects: abstract_object(i, mapper, sregistry) - # from IPython import embed; embed() return mapper @@ -455,6 +454,7 @@ def _(i, mapper, sregistry): @abstract_object.register(Array) +@abstract_object.register(PETScArray) def _(i, mapper, sregistry): if isinstance(i, Lock): name = sregistry.make_name(prefix='lock') @@ -472,20 +472,6 @@ def _(i, mapper, sregistry): mapper[i.dmap] = v.dmap -# @abstract_object.register(PETScArray) -# def _(i, mapper, sregistry): -# name = sregistry.make_name(prefix='xx') - -# v = i._rebuild(name=name, initializer=None, alias=True) - -# mapper.update({ -# i: v, -# i.indexed: v.indexed, -# i.dmap: v.dmap, -# i._C_symbol: v._C_symbol, -# }) - - @abstract_object.register(Bundle) def _(i, mapper, sregistry): name = sregistry.make_name(prefix='a') @@ -508,7 +494,6 @@ def _(i, mapper, sregistry): name = sregistry.make_name(prefix='o') v = i._rebuild(name) - mapper[i] = v diff --git a/devito/petsc/types/object.py b/devito/petsc/types/object.py index 7821e316d9..a285863d14 100644 --- a/devito/petsc/types/object.py +++ b/devito/petsc/types/object.py @@ -218,9 +218,6 @@ def __init__(self, name='subctx', pname='SubMatrixCtx', fields=None, _C_modifier = None -JacobianStructCast = cast('struct JacobianCtx *') - - class PETScArrayObject(ArrayObject): _data_alignment = False diff --git a/devito/types/object.py b/devito/types/object.py index 6eea49f6a1..1153ce81eb 100644 --- a/devito/types/object.py +++ b/devito/types/object.py @@ -249,8 +249,10 @@ class LocalCompositeObject(CompositeObject, LocalType): """ __rargs__ = ('name', 'pname', 'fields') + __rkwargs__ = ('modifier', 'liveness') def __init__(self, name, pname, fields, modifier=None, liveness='lazy'): + self.modifier = modifier dtype = CustomDtype(f"struct {pname}", modifier=modifier) Object.__init__(self, name, dtype, None) self._pname = pname From 68346110438eadb1d43e27786f22fe1f333ddbca Mon Sep 17 00:00:00 2001 From: ZoeLeibowitz Date: Thu, 29 May 2025 18:16:54 +0100 Subject: [PATCH 29/84] dsl/compiler: Zero columns for essential bcs --- devito/passes/iet/engine.py | 6 +-- devito/petsc/solve.py | 45 ++++++++++++++++--- devito/petsc/types/types.py | 3 -- devito/tools/utils.py | 2 +- examples/petsc/Poisson/01_poisson.py | 22 +++++---- examples/petsc/Poisson/02_laplace.py | 6 ++- examples/petsc/Poisson/04_poisson.py | 4 +- .../petsc/random/biharmonic/02_biharmonic.py | 18 +++----- tests/test_petsc.py | 16 ++++--- 9 files changed, 79 insertions(+), 43 deletions(-) diff --git a/devito/passes/iet/engine.py b/devito/passes/iet/engine.py index 0c3c7481bf..fa34a4946c 100644 --- a/devito/passes/iet/engine.py +++ b/devito/passes/iet/engine.py @@ -416,8 +416,7 @@ def abstract_objects(objects0, sregistry=None): # Precedence rules make it possible to reconstruct objects that depend on # higher priority objects - keys = [Bundle, Array, PETScArray, DiscreteFunction, AbstractIncrDimension, BlockDimension] - # keys = [Bundle, Array, DiscreteFunction, AbstractIncrDimension, BlockDimension] + keys = [Bundle, ArrayBasic, DiscreteFunction, AbstractIncrDimension, BlockDimension] priority = {k: i for i, k in enumerate(keys, start=1)} objects = sorted_priority(objects, priority) @@ -453,8 +452,7 @@ def _(i, mapper, sregistry): }) -@abstract_object.register(Array) -@abstract_object.register(PETScArray) +@abstract_object.register(ArrayBasic) def _(i, mapper, sregistry): if isinstance(i, Lock): name = sregistry.make_name(prefix='lock') diff --git a/devito/petsc/solve.py b/devito/petsc/solve.py index c34b829798..f5893fec05 100644 --- a/devito/petsc/solve.py +++ b/devito/petsc/solve.py @@ -52,10 +52,18 @@ def linear_solve_args(self): return target, tuple(funcs), self.generate_field_data(eqns, target, arrays) def generate_field_data(self, eqns, target, arrays): - formfuncs, formrhs = zip( - *[self.build_function_eqns(eq, target, arrays) for eq in eqns] + # TODO: Ensure essential BCs are handled first - this is required to + # maintain the symmetry of the operator when "constructing" the Jacobian. + eqns = sorted(eqns, key=lambda e: 0 if isinstance(e, EssentialBC) else 1) + + # TODO: scaling + + formfuncs, formrhs = map( + lambda x: [e for i in x for e in (i if isinstance(i, tuple) else [i])], + zip(*[self.build_function_eqns(eq, target, arrays) for eq in eqns]) ) - matvecs = [self.build_matvec_eqns(eq, target, arrays) for eq in eqns] + matvecs = [e for i in [self.build_matvec_eqns(eq, target, arrays) for eq in eqns] + for e in (i if isinstance(i, (tuple)) else [i])] initialguess = [ eq for eq in @@ -88,7 +96,14 @@ def build_matvec_eqns(self, eq, target, arrays): def make_matvec(self, eq, F_target, arrays, targets): if isinstance(eq, EssentialBC): + # TODO: SCALING + # NOTE: Until PetscSection + DMDA is supported, we leave + # the essential BCs in the solver. + # Trivial equations for bc rows -> place 1.0 on diagonal (scaled) + # and zero symmetrically. rhs = arrays['x'] + zero_column = Eq(arrays['x'], 0.0, subdomain=eq.subdomain) + return (Eq(arrays['y'], rhs, subdomain=eq.subdomain), zero_column) else: rhs = F_target.subs(targets_to_arrays(arrays['x'], targets)) rhs = rhs.subs(self.time_mapper) @@ -96,9 +111,16 @@ def make_matvec(self, eq, F_target, arrays, targets): def make_formfunc(self, eq, F_target, arrays, targets): if isinstance(eq, EssentialBC): - # TODO: CHECK THIS - rhs = arrays[main_target]['x'] - eq.rhs - # rhs = 0. + # The initial guess already satisfies the essential boundary conditions, + # so this term will always be zero. It's included here to allow + # testing of the Jacobian via finite differences. + rhs = arrays['x'] - eq.rhs + # Create equation to handle essential boundary conditions - we + # move the essential BCs to the right-hand side + # and zero the corresponding column in the Jacobian. + # TODO: extend this to mixed problems + move_bc = Eq(arrays['x'], eq.rhs, subdomain=eq.subdomain) + return (Eq(arrays['f'], rhs, subdomain=eq.subdomain), move_bc) else: if isinstance(F_target, (int, float)): rhs = F_target @@ -215,6 +237,17 @@ def generate_arrays_combined(self, *targets): class EssentialBC(Eq): + """ + A special equation representing an essential boundary condition. + It is used to handle essential boundary conditions in the PETSc solver. + Until PetscSection + DMDA is supported, we treat essential boundary conditions + as trivial equations in the solver, where we place 1.0 on the diagonal of the jacobian, + zero symmetrically and move the essential boundary condition to the right-hand side. + + NOTE: When users define essential boundary conditions, they need to ensure that + the SubDomains do not overlap. Solver will still run but may see unexpected behaviour + along boundaries. + """ pass diff --git a/devito/petsc/types/types.py b/devito/petsc/types/types.py index 973d558f38..afab9fff63 100644 --- a/devito/petsc/types/types.py +++ b/devito/petsc/types/types.py @@ -188,7 +188,6 @@ def targets(self): return (self.target,) -# TODO: should this acc inherhit from fielddata? maybe not? class MultipleFieldData(FieldData): def __init__(self, targets, arrays, submatrices=None): self._targets = as_tuple(targets) @@ -196,8 +195,6 @@ def __init__(self, targets, arrays, submatrices=None): self._submatrices = submatrices self._formfuncs = [] - # pass - def extend_formfuncs(self, formfuncs): self._formfuncs.extend(formfuncs) diff --git a/devito/tools/utils.py b/devito/tools/utils.py index 0657146040..0920298356 100644 --- a/devito/tools/utils.py +++ b/devito/tools/utils.py @@ -337,7 +337,7 @@ def sorted_priority(items, priority): """ def key(i): - for cls in sorted(priority, key=priority.get, reverse=True): + for cls in sorted(priority, key=priority.get): if isinstance(i, cls): v = priority[cls] break diff --git a/examples/petsc/Poisson/01_poisson.py b/examples/petsc/Poisson/01_poisson.py index 6d89091cd6..86fdfc6dd3 100644 --- a/examples/petsc/Poisson/01_poisson.py +++ b/examples/petsc/Poisson/01_poisson.py @@ -16,6 +16,8 @@ # Subdomains to implement BCs +# NOTE: For essential BCs, we ensure the SubDomains do not overlap + class SubTop(SubDomain): name = 'subtop' @@ -37,7 +39,7 @@ class SubLeft(SubDomain): def define(self, dimensions): x, y = dimensions - return {x: ('left', 1), y: y} + return {x: ('left', 1), y: ('middle', 1, 1)} class SubRight(SubDomain): @@ -45,7 +47,7 @@ class SubRight(SubDomain): def define(self, dimensions): x, y = dimensions - return {x: ('right', 1), y: y} + return {x: ('right', 1), y: ('middle', 1, 1)} sub1 = SubTop() @@ -64,6 +66,7 @@ def analytical(x, y): Ly = np.float64(1.) n_values = list(range(13, 174, 10)) +n_values = [6] dx = np.array([Lx/(n-1) for n in n_values]) errors = [] @@ -76,6 +79,7 @@ def analytical(x, y): phi = Function(name='phi', grid=grid, space_order=2, dtype=np.float64) rhs = Function(name='rhs', grid=grid, space_order=2, dtype=np.float64) + bc = Function(name='bc', grid=grid, space_order=2, dtype=np.float64) eqn = Eq(rhs, phi.laplace, subdomain=grid.interior) @@ -87,11 +91,13 @@ def analytical(x, y): 2.0*X*(Y-1.0)*(Y - 2.0*X + X*Y + 2.0) ) * np.float64(np.exp(X-Y)) + bc.data[:] = 0.0 + # # Create boundary condition expressions using subdomains - bcs = [EssentialBC(phi, np.float64(0.), subdomain=sub1)] - bcs += [EssentialBC(phi, np.float64(0.), subdomain=sub2)] - bcs += [EssentialBC(phi, np.float64(0.), subdomain=sub3)] - bcs += [EssentialBC(phi, np.float64(0.), subdomain=sub4)] + bcs = [EssentialBC(phi, bc, subdomain=sub1)] + bcs += [EssentialBC(phi, bc, subdomain=sub2)] + bcs += [EssentialBC(phi, bc, subdomain=sub3)] + bcs += [EssentialBC(phi, bc, subdomain=sub4)] exprs = [eqn] + bcs petsc = PETScSolve(exprs, target=phi, solver_parameters={'ksp_rtol': 1e-8}) @@ -108,5 +114,5 @@ def analytical(x, y): slope, _ = np.polyfit(np.log(dx), np.log(errors), 1) -assert slope > 1.9 -assert slope < 2.1 +# assert slope > 1.9 +# assert slope < 2.1 diff --git a/examples/petsc/Poisson/02_laplace.py b/examples/petsc/Poisson/02_laplace.py index 780da10ec5..c84855137f 100644 --- a/examples/petsc/Poisson/02_laplace.py +++ b/examples/petsc/Poisson/02_laplace.py @@ -23,12 +23,14 @@ # Subdomains to implement BCs +# NOTE: For essential BCs, we ensure the SubDomains do not overlap + class SubTop(SubDomain): name = 'subtop' def define(self, dimensions): x, y = dimensions - return {x: x, y: ('right', 1)} + return {x: ('middle', 1, 1), y: ('right', 1)} class SubBottom(SubDomain): @@ -36,7 +38,7 @@ class SubBottom(SubDomain): def define(self, dimensions): x, y = dimensions - return {x: x, y: ('left', 1)} + return {x: ('middle', 1, 1), y: ('left', 1)} class SubLeft(SubDomain): diff --git a/examples/petsc/Poisson/04_poisson.py b/examples/petsc/Poisson/04_poisson.py index 44f34ec8f9..637ce44076 100644 --- a/examples/petsc/Poisson/04_poisson.py +++ b/examples/petsc/Poisson/04_poisson.py @@ -26,7 +26,7 @@ class SubTop(SubDomain): def define(self, dimensions): x, y = dimensions - return {x: x, y: ('right', 1)} + return {x: ('middle', 1, 1), y: ('right', 1)} class SubBottom(SubDomain): @@ -34,7 +34,7 @@ class SubBottom(SubDomain): def define(self, dimensions): x, y = dimensions - return {x: x, y: ('left', 1)} + return {x: ('middle', 1, 1), y: ('left', 1)} class SubLeft(SubDomain): diff --git a/examples/petsc/random/biharmonic/02_biharmonic.py b/examples/petsc/random/biharmonic/02_biharmonic.py index 8c15f3d3a2..6db85d6d96 100644 --- a/examples/petsc/random/biharmonic/02_biharmonic.py +++ b/examples/petsc/random/biharmonic/02_biharmonic.py @@ -76,8 +76,7 @@ def f_fcn(x, y): Lx = np.float64(1.) Ly = np.float64(1.) -# n_values = [33, 53, 73, 93, 113] -n_values = [33] +n_values = [33, 53, 73, 93, 113] dx = np.array([Lx/(n-1) for n in n_values]) u_errors = [] @@ -135,14 +134,11 @@ def f_fcn(x, y): v_error = np.linalg.norm(v_diff.ravel(), ord=np.inf) / np.linalg.norm(lap_u.data[:].ravel(), ord=np.inf) v_errors.append(v_error) -# u_slope, _ = np.polyfit(np.log(dx), np.log(u_errors), 1) -# v_slope, _ = np.polyfit(np.log(dx), np.log(v_errors), 1) +u_slope, _ = np.polyfit(np.log(dx), np.log(u_errors), 1) +v_slope, _ = np.polyfit(np.log(dx), np.log(v_errors), 1) -# assert u_slope > 1.9 -# assert u_slope < 2.1 +assert u_slope > 1.9 +assert u_slope < 2.1 -# assert v_slope > 1.9 -# assert v_slope < 2.1 -print(op.ccode) -print("u_errors:", u_errors) -print("v_errors:", v_errors) +assert v_slope > 1.9 +assert v_slope < 2.1 diff --git a/tests/test_petsc.py b/tests/test_petsc.py index 943331f3c3..5c1346287f 100644 --- a/tests/test_petsc.py +++ b/tests/test_petsc.py @@ -339,7 +339,7 @@ def test_separate_eqn(eqn, target, expected): f2 = TimeFunction(name='f2', grid=grid, space_order=so) # noqa g2 = TimeFunction(name='g2', grid=grid, space_order=so) # noqa - b, F, _ = separate_eqn(eval(eqn), eval(target)) + b, F, _, _= separate_eqn(eval(eqn), eval(target)) expected_b, expected_F = expected assert str(b) == expected_b @@ -471,7 +471,7 @@ def test_separate_eval_eqn(eqn, target, expected): f2 = TimeFunction(name='f2', grid=grid, space_order=so) # noqa g2 = TimeFunction(name='g2', grid=grid, space_order=so) # noqa - b, F, _ = separate_eqn(eval(eqn), eval(target)) + b, F, _, _ = separate_eqn(eval(eqn), eval(target)) expected_b, expected_F = expected assert str(b) == expected_b @@ -793,6 +793,9 @@ def test_essential_bcs(): Verify that PETScSolve returns the correct output with essential boundary conditions. """ + + # SubDomains used for essential boundary conditions + # should not overlap. class SubTop(SubDomain): name = 'subtop' @@ -814,7 +817,7 @@ class SubLeft(SubDomain): def define(self, dimensions): x, y = dimensions - return {x: ('left', 1), y: y} + return {x: ('left', 1), y: ('middle', 1, 1)} sub3 = SubLeft() class SubRight(SubDomain): @@ -822,7 +825,7 @@ class SubRight(SubDomain): def define(self, dimensions): x, y = dimensions - return {x: ('right', 1), y: y} + return {x: ('right', 1), y: ('middle', 1, 1)} sub4 = SubRight() subdomains = (sub1, sub2, sub3, sub4) @@ -913,9 +916,10 @@ def test_coupled_vs_non_coupled(self): callbacks2 = [meta_call.root for meta_call in op2._func_table.values()] # Solving for multiple fields within the same matrix system requires - # additional machinery and more callback functions + # less callback functions than solving them separately. + # TODO: check reuse of callback functions where appropriate assert len(callbacks1) == 8 - assert len(callbacks2) == 11 + assert len(callbacks2) == 6 # Check fielddata type fielddata1 = petsc1[0].rhs.fielddata From 9642278b4c19ec0c0adfc1a2d0ba0208569e5b37 Mon Sep 17 00:00:00 2001 From: ZoeLeibowitz Date: Thu, 29 May 2025 18:51:52 +0100 Subject: [PATCH 30/84] dsl: Compatible scaling of jacobian --- devito/petsc/solve.py | 52 +++++++++++++------ examples/petsc/Poisson/01_poisson.py | 6 +-- .../petsc/random/biharmonic/02_biharmonic.py | 4 +- .../random/biharmonic/biharmonic_matfree.c | 2 +- tests/test_petsc.py | 13 ++--- 5 files changed, 49 insertions(+), 28 deletions(-) diff --git a/devito/petsc/solve.py b/devito/petsc/solve.py index f5893fec05..d675fb0cb9 100644 --- a/devito/petsc/solve.py +++ b/devito/petsc/solve.py @@ -1,6 +1,7 @@ from functools import singledispatch import sympy +import numpy as np from devito.finite_differences.differentiable import Mul from devito.finite_differences.derivative import Derivative @@ -28,9 +29,11 @@ def __init__(self, solver_parameters=None, target_eqns=None): self.solver_params = solver_parameters self.time_mapper = None self.target_eqns = target_eqns + self.cell_area = None def build_eq(self): target, funcs, fielddata = self.linear_solve_args() + # Placeholder equation for inserting calls to the solver linear_solve = LinearSolveExpr( funcs, @@ -44,6 +47,7 @@ def build_eq(self): def linear_solve_args(self): target, eqns = next(iter(self.target_eqns.items())) eqns = as_tuple(eqns) + self.cell_area = np.prod(target.grid.spacing_symbols) funcs = get_funcs(eqns) self.time_mapper = generate_time_mapper(funcs) @@ -52,19 +56,17 @@ def linear_solve_args(self): return target, tuple(funcs), self.generate_field_data(eqns, target, arrays) def generate_field_data(self, eqns, target, arrays): - # TODO: Ensure essential BCs are handled first - this is required to - # maintain the symmetry of the operator when "constructing" the Jacobian. + # TODO: Apply essential boundary conditions first to preserve operator symmetry + # during Jacobian "construction". eqns = sorted(eqns, key=lambda e: 0 if isinstance(e, EssentialBC) else 1) - # TODO: scaling - formfuncs, formrhs = map( lambda x: [e for i in x for e in (i if isinstance(i, tuple) else [i])], zip(*[self.build_function_eqns(eq, target, arrays) for eq in eqns]) ) matvecs = [e for i in [self.build_matvec_eqns(eq, target, arrays) for eq in eqns] for e in (i if isinstance(i, (tuple)) else [i])] - + from IPython import embed; embed() initialguess = [ eq for eq in (self.make_initial_guess(e, target, arrays) for e in eqns) @@ -96,17 +98,17 @@ def build_matvec_eqns(self, eq, target, arrays): def make_matvec(self, eq, F_target, arrays, targets): if isinstance(eq, EssentialBC): - # TODO: SCALING # NOTE: Until PetscSection + DMDA is supported, we leave # the essential BCs in the solver. # Trivial equations for bc rows -> place 1.0 on diagonal (scaled) # and zero symmetrically. rhs = arrays['x'] zero_column = Eq(arrays['x'], 0.0, subdomain=eq.subdomain) - return (Eq(arrays['y'], rhs, subdomain=eq.subdomain), zero_column) + return (EssentialBC(arrays['y'], rhs, subdomain=eq.subdomain), zero_column) else: rhs = F_target.subs(targets_to_arrays(arrays['x'], targets)) rhs = rhs.subs(self.time_mapper) + rhs = rhs * self.cell_area return Eq(arrays['y'], rhs, subdomain=eq.subdomain) def make_formfunc(self, eq, F_target, arrays, targets): @@ -120,17 +122,19 @@ def make_formfunc(self, eq, F_target, arrays, targets): # and zero the corresponding column in the Jacobian. # TODO: extend this to mixed problems move_bc = Eq(arrays['x'], eq.rhs, subdomain=eq.subdomain) - return (Eq(arrays['f'], rhs, subdomain=eq.subdomain), move_bc) + return (EssentialBC(arrays['f'], rhs, subdomain=eq.subdomain), move_bc) else: if isinstance(F_target, (int, float)): - rhs = F_target + rhs = F_target * self.cell_area else: rhs = F_target.subs(targets_to_arrays(arrays['x'], targets)) rhs = rhs.subs(self.time_mapper) + rhs = rhs * self.cell_area return Eq(arrays['f'], rhs, subdomain=eq.subdomain) def make_rhs(self, eq, b, arrays): rhs = 0. if isinstance(eq, EssentialBC) else b.subs(self.time_mapper) + rhs = rhs*self.cell_area return Eq(arrays['b'], rhs, subdomain=eq.subdomain) def make_initial_guess(self, eq, target, arrays): @@ -175,6 +179,8 @@ def linear_solve_args(self): all_data = MultipleFieldData(submatrices=jacobian, arrays=arrays, targets=coupled_targets) + self.cell_area = np.prod(all_data.grid.spacing_symbols) + for target, eqns in self.target_eqns.items(): eqns = as_tuple(eqns) self.update_jacobian(eqns, target, jacobian, arrays[target]) @@ -189,11 +195,11 @@ def linear_solve_args(self): def update_jacobian(self, eqns, target, jacobian, arrays): for submat, mtvs in jacobian.submatrices[target].items(): - matvecs = [ - self.build_matvec_eqns(eq, mtvs['derivative_wrt'], arrays) - for eq in eqns - ] + # TODO: maintain symmetry for coupled + matvecs = [e for i in [self.build_matvec_eqns(eq, mtvs['derivative_wrt'], arrays) + for eq in eqns] for e in (i if isinstance(i, (tuple)) else [i])] matvecs = [m for m in matvecs if m is not None] + if matvecs: jacobian.set_submatrix(target, submat, matvecs) @@ -210,14 +216,14 @@ def build_function_eqns(self, eq, main_target, coupled_targets, arrays): mapper.update(targets_to_arrays(arrays[t]['x'], target_funcs)) if isinstance(eq, EssentialBC): - # TODO: CHECK THIS rhs = arrays[main_target]['x'] - eq.rhs else: if isinstance(zeroed, (int, float)): - rhs = zeroed + rhs = zeroed * self.cell_area else: rhs = zeroed.subs(mapper) rhs = rhs.subs(self.time_mapper) + rhs = rhs * self.cell_area return Eq(arrays[main_target]['f'], rhs, subdomain=eq.subdomain) @@ -235,6 +241,20 @@ def generate_arrays_combined(self, *targets): for target in targets } + def make_matvec(self, eq, F_target, arrays, targets): + if isinstance(eq, EssentialBC): + # TODO: SHOULDN'T NEED this subclass once fixed mixed problems + symmetry + # NOTE: Until PetscSection + DMDA is supported, we leave + # the essential BCs in the solver. + # Trivial equations for bc rows -> place 1.0 on diagonal + # and zero symmetrically. FIX FOR MIXED PROBLEMS + rhs = arrays['x'] + else: + rhs = F_target.subs(targets_to_arrays(arrays['x'], targets)) + rhs = rhs.subs(self.time_mapper) + rhs = rhs * self.cell_area + return Eq(arrays['y'], rhs, subdomain=eq.subdomain) + class EssentialBC(Eq): """ @@ -246,7 +266,7 @@ class EssentialBC(Eq): NOTE: When users define essential boundary conditions, they need to ensure that the SubDomains do not overlap. Solver will still run but may see unexpected behaviour - along boundaries. + at boundaries. """ pass diff --git a/examples/petsc/Poisson/01_poisson.py b/examples/petsc/Poisson/01_poisson.py index 86fdfc6dd3..6f24f74f64 100644 --- a/examples/petsc/Poisson/01_poisson.py +++ b/examples/petsc/Poisson/01_poisson.py @@ -66,7 +66,7 @@ def analytical(x, y): Ly = np.float64(1.) n_values = list(range(13, 174, 10)) -n_values = [6] +n_values = [9] dx = np.array([Lx/(n-1) for n in n_values]) errors = [] @@ -114,5 +114,5 @@ def analytical(x, y): slope, _ = np.polyfit(np.log(dx), np.log(errors), 1) -# assert slope > 1.9 -# assert slope < 2.1 +assert slope > 1.9 +assert slope < 2.1 diff --git a/examples/petsc/random/biharmonic/02_biharmonic.py b/examples/petsc/random/biharmonic/02_biharmonic.py index 6db85d6d96..bc7b54654d 100644 --- a/examples/petsc/random/biharmonic/02_biharmonic.py +++ b/examples/petsc/random/biharmonic/02_biharmonic.py @@ -20,7 +20,7 @@ class SubTop(SubDomain): def define(self, dimensions): x, y = dimensions - return {x: x, y: ('right', 1)} + return {x: ('middle', 1, 1), y: ('right', 1)} class SubBottom(SubDomain): @@ -28,7 +28,7 @@ class SubBottom(SubDomain): def define(self, dimensions): x, y = dimensions - return {x: x, y: ('left', 1)} + return {x: ('middle', 1, 1), y: ('left', 1)} class SubLeft(SubDomain): diff --git a/examples/petsc/random/biharmonic/biharmonic_matfree.c b/examples/petsc/random/biharmonic/biharmonic_matfree.c index 53121473e7..2fb432382f 100644 --- a/examples/petsc/random/biharmonic/biharmonic_matfree.c +++ b/examples/petsc/random/biharmonic/biharmonic_matfree.c @@ -97,7 +97,7 @@ int main(int argc,char **argv) { user.f = &f_fcn; PetscCall(DMDACreate2d(PETSC_COMM_WORLD, DM_BOUNDARY_NONE, DM_BOUNDARY_NONE, DMDA_STENCIL_STAR, - 33,33,PETSC_DECIDE,PETSC_DECIDE, + 6,6,PETSC_DECIDE,PETSC_DECIDE, 2,1, // degrees of freedom, stencil width NULL,NULL,&da)); PetscCall(DMSetApplicationContext(da,&user)); diff --git a/tests/test_petsc.py b/tests/test_petsc.py index 5c1346287f..c136160ecc 100644 --- a/tests/test_petsc.py +++ b/tests/test_petsc.py @@ -108,15 +108,15 @@ def test_petsc_solve(): rhs_expr = FindNodes(Expression).visit(formrhs_callback[0]) assert str(action_expr[-1].expr.rhs) == ( - 'x_f[x + 1, y + 2]/ctx0->h_x**2' + '(x_f[x + 1, y + 2]/ctx0->h_x**2' ' - 2.0*x_f[x + 2, y + 2]/ctx0->h_x**2' ' + x_f[x + 3, y + 2]/ctx0->h_x**2' ' + x_f[x + 2, y + 1]/ctx0->h_y**2' ' - 2.0*x_f[x + 2, y + 2]/ctx0->h_y**2' - ' + x_f[x + 2, y + 3]/ctx0->h_y**2' + ' + x_f[x + 2, y + 3]/ctx0->h_y**2)*ctx0->h_x*ctx0->h_y' ) - assert str(rhs_expr[-1].expr.rhs) == 'g[x + 2, y + 2]' + assert str(rhs_expr[-1].expr.rhs) == 'ctx0->h_x*ctx0->h_y*g[x + 2, y + 2]' # Check the iteration bounds are correct. assert op.arguments().get('x_m') == 0 @@ -836,7 +836,7 @@ def define(self, dimensions): # Solving Ax=b where A is the identity matrix v.data[:] = 5.0 - eqn = Eq(u, v) + eqn = Eq(u, v, subdomain=grid.interior) bcs = [EssentialBC(u, 1., subdomain=sub1)] # top bcs += [EssentialBC(u, 2., subdomain=sub2)] # bottom @@ -1083,11 +1083,12 @@ def test_submatrices(self): j00 = submatrices.get_submatrix(e, 'J00') j11 = submatrices.get_submatrix(g, 'J11') + # Compatible scaling to reduce condition number of jacobian assert str(j00['matvecs'][0]) == 'Eq(y_e(x, y),' \ - + ' Derivative(x_e(x, y), (x, 2)) + Derivative(x_e(x, y), (y, 2)))' + + ' h_x*h_y*(Derivative(x_e(x, y), (x, 2)) + Derivative(x_e(x, y), (y, 2))))' assert str(j11['matvecs'][0]) == 'Eq(y_g(x, y),' \ - + ' Derivative(x_g(x, y), (x, 2)) + Derivative(x_g(x, y), (y, 2)))' + + ' h_x*h_y*(Derivative(x_g(x, y), (x, 2)) + Derivative(x_g(x, y), (y, 2))))' # Check the derivative wrt fields assert j00['derivative_wrt'] == e From ce9150f71a2d07513b6f9e9d6787a7e4ab004ffa Mon Sep 17 00:00:00 2001 From: ZoeLeibowitz Date: Fri, 30 May 2025 17:38:20 +0100 Subject: [PATCH 31/84] dsl: Scale boundary rows of jacobian --- devito/passes/iet/engine.py | 7 +- devito/petsc/iet/routines.py | 44 +-- devito/petsc/solve.py | 255 ++++++++++++------ devito/petsc/types/array.py | 13 +- devito/petsc/types/types.py | 12 + devito/tools/utils.py | 1 - examples/petsc/Poisson/01_poisson.py | 8 +- .../petsc/random/biharmonic/02_biharmonic.py | 18 +- tests/test_petsc.py | 111 +++++++- 9 files changed, 343 insertions(+), 126 deletions(-) diff --git a/devito/passes/iet/engine.py b/devito/passes/iet/engine.py index fa34a4946c..ad0ebefd1e 100644 --- a/devito/passes/iet/engine.py +++ b/devito/passes/iet/engine.py @@ -10,14 +10,13 @@ from devito.passes import needs_transfer from devito.symbolics import FieldFromComposite, FieldFromPointer from devito.tools import DAG, as_tuple, filter_ordered, sorted_priority, timed_pass -from devito.types import (Array, Bundle, CompositeObject, Lock, IncrDimension, +from devito.types import (Bundle, CompositeObject, Lock, IncrDimension, ModuloDimension, Indirection, Pointer, SharedData, ThreadArray, Temp, NPThreads, NThreadsBase, Wildcard) from devito.types.array import ArrayBasic from devito.types.args import ArgProvider from devito.types.dense import DiscreteFunction from devito.types.dimension import AbstractIncrDimension, BlockDimension -from devito.petsc.types import PETScArray __all__ = ['Graph', 'iet_pass', 'iet_visit'] @@ -150,7 +149,6 @@ def apply(self, func, **kwargs): # Minimize code size if len(efuncs) > len(self.efuncs): efuncs = reuse_compounds(efuncs, self.sregistry) - # TODO: fix for petsc bundles efuncs = reuse_efuncs(self.root, efuncs, self.sregistry) self.efuncs = efuncs @@ -423,7 +421,6 @@ def abstract_objects(objects0, sregistry=None): # Build abstraction mappings mapper = {} sregistry = sregistry or SymbolRegistry() - for i in objects: abstract_object(i, mapper, sregistry) @@ -473,7 +470,6 @@ def _(i, mapper, sregistry): @abstract_object.register(Bundle) def _(i, mapper, sregistry): name = sregistry.make_name(prefix='a') - components = [mapper[f] for f in i.components] v = i._rebuild(name=name, components=components, alias=True) @@ -492,6 +488,7 @@ def _(i, mapper, sregistry): name = sregistry.make_name(prefix='o') v = i._rebuild(name) + mapper[i] = v diff --git a/devito/petsc/iet/routines.py b/devito/petsc/iet/routines.py index 8674dc4df4..7c824ad638 100644 --- a/devito/petsc/iet/routines.py +++ b/devito/petsc/iet/routines.py @@ -793,7 +793,7 @@ def _whole_formfunc_body(self, body, fielddata): bundles = sobjs['bundles'] fbundle = bundles['f'] xbundle = bundles['x'] - body = self.bundle_residual(body, bundles) + body = self.residual_bundle(body, bundles) dm_cast = DummyExpr(dmda, DMCast(objs['dummyptr']), init=True) @@ -1046,13 +1046,9 @@ def _submat_callback_body(self): stacks=(get_ctx, deref_subdm), retstmt=(Call('PetscFunctionReturn', arguments=[0]),) ) - - def bundle_residual(self, body, bundles): - fbundle = bundles['f'] - xbundle = bundles['x'] + def residual_bundle(self, body, bundles): mapper1 = bundles['bundle_mapper'] - indexeds = FindSymbols('indexeds').visit(body) subss = {} @@ -1067,7 +1063,6 @@ def bundle_residual(self, body, bundles): return body - class BaseObjectBuilder: """ A base class for constructing objects needed for a PETSc solver. @@ -1180,7 +1175,6 @@ def _extend_build(self, base_dict): base_dict[f'{key}Y'] = CallbackVec(f'{key}Y') base_dict[f'{key}F'] = CallbackVec(f'{key}F') - # Bundle objects/metadata required by the coupled residual callback f_components = [] x_components = [] @@ -1194,9 +1188,13 @@ def _extend_build(self, base_dict): f_components.append(f_arr) x_components.append(x_arr) - # TODO: to group them, maybe pass in struct name as arg to petscbundle - fbundle = PetscBundle(name='f_bundle', components=f_components) - xbundle = PetscBundle(name='x_bundle', components=x_components) + bundle_pname = sreg.make_name(prefix='Field') + fbundle = PetscBundle( + name='f_bundle', components=f_components, pname=bundle_pname + ) + xbundle = PetscBundle( + name='x_bundle', components=x_components, pname=bundle_pname + ) # Build the bundle mapper for i, t in enumerate(targets): @@ -1205,7 +1203,6 @@ def _extend_build(self, base_dict): bundle_mapper[f_arr.base] = fbundle bundle_mapper[x_arr.base] = xbundle - base_dict['bundles'] = { 'f': fbundle, 'x': xbundle, @@ -1346,6 +1343,10 @@ def _setup(self): self.snes_ctx] ) + snes_set_options = petsc_call( + 'SNESSetFromOptions', [sobjs['snes']] + ) + dmda_calls = self._create_dmda_calls(dmda) mainctx = sobjs['userctx'] @@ -1377,6 +1378,7 @@ def _setup(self): ksp_set_from_ops, matvec_operation, formfunc_operation, + snes_set_options, call_struct_callback, mat_set_dm, calls_set_app_ctx, @@ -1504,6 +1506,10 @@ def _setup(self): self.snes_ctx] ) + snes_set_options = petsc_call( + 'SNESSetFromOptions', [sobjs['snes']] + ) + dmda_calls = self._create_dmda_calls(dmda) mainctx = sobjs['userctx'] @@ -1574,6 +1580,7 @@ def _setup(self): ksp_set_from_ops, matvec_operation, formfunc_operation, + snes_set_options, call_struct_callback, mat_set_dm, calls_set_app_ctx, @@ -1670,17 +1677,12 @@ def _execute_solve(self): """ objs = self.objs sobjs = self.solver_objs - - struct_assignment = self.timedep.assign_time_iters(sobjs['userctx']) - - # rhs_callbacks = self.cbbuilder.formrhs - xglob = sobjs['xglobal'] - bglob = sobjs['bglobal'] + struct_assignment = self.timedep.assign_time_iters(sobjs['userctx']) targets = self.injectsolve.expr.rhs.fielddata.targets - # TODO: optimise the ccode generated here + # TODO: Optimise the ccode generated here pre_solve = () post_solve = () @@ -1693,10 +1695,10 @@ def _execute_solve(self): s = sobjs[f'scatter{name}'] pre_solve += ( - # TODO: switch to createwitharray and move to setup + # TODO: Switch to createwitharray and move to setup petsc_call('DMCreateLocalVector', [dm, Byref(target_xloc)]), - # TODO: need to call reset array + # TODO: Need to call reset array self.timedep.place_array(t), petsc_call( 'DMLocalToGlobal', diff --git a/devito/petsc/solve.py b/devito/petsc/solve.py index d675fb0cb9..278a5092b8 100644 --- a/devito/petsc/solve.py +++ b/devito/petsc/solve.py @@ -2,6 +2,8 @@ import sympy import numpy as np +from itertools import chain +from collections import defaultdict from devito.finite_differences.differentiable import Mul from devito.finite_differences.derivative import Derivative @@ -29,7 +31,14 @@ def __init__(self, solver_parameters=None, target_eqns=None): self.solver_params = solver_parameters self.time_mapper = None self.target_eqns = target_eqns + # TODO: make this _ self.cell_area = None + # self._centre_stencils = set() + self._diag_scale = defaultdict(set) + + @property + def diag_scale(self): + return self._diag_scale def build_eq(self): target, funcs, fielddata = self.linear_solve_args() @@ -56,17 +65,35 @@ def linear_solve_args(self): return target, tuple(funcs), self.generate_field_data(eqns, target, arrays) def generate_field_data(self, eqns, target, arrays): - # TODO: Apply essential boundary conditions first to preserve operator symmetry - # during Jacobian "construction". + # Apply essential boundary conditions first to preserve + # operator symmetry during Jacobian "construction" eqns = sorted(eqns, key=lambda e: 0 if isinstance(e, EssentialBC) else 1) + matvecs = [e for eq in eqns for e in self.build_matvec_eq(eq, target, arrays)] + formfuncs, formrhs = map( - lambda x: [e for i in x for e in (i if isinstance(i, tuple) else [i])], - zip(*[self.build_function_eqns(eq, target, arrays) for eq in eqns]) + lambda x: [e for i in x for e in i], + zip(*[self.build_function_eq(eq, target, arrays) for eq in eqns]) ) - matvecs = [e for i in [self.build_matvec_eqns(eq, target, arrays) for eq in eqns] - for e in (i if isinstance(i, (tuple)) else [i])] - from IPython import embed; embed() + + # self._centre_stencils[arrays['x']].update( + stencils = set() + for eq in matvecs: + if not isinstance(eq, EssentialBC): + stencil = centre_stencil(eq.rhs, arrays['x'], as_coeff=True) + stencils.add(stencil) + + if len(stencils) > 1: + # Scaling of jacobian is therefore ambiguous, potentially could average across the subblock + # for now just set to trivial 1.0 + scale = 1.0 + else: + scale = next(iter(stencils)) + + # from IPython import embed; embed() + matvecs = self.scale_essential_bcs(matvecs, scale) + formfuncs = self.scale_essential_bcs(formfuncs, scale) + initialguess = [ eq for eq in (self.make_initial_guess(e, target, arrays) for e in eqns) @@ -82,19 +109,18 @@ def generate_field_data(self, eqns, target, arrays): arrays=arrays ) - def build_function_eqns(self, eq, target, arrays): + def build_function_eq(self, eq, target, arrays): b, F_target, _, targets = separate_eqn(eq, target) formfunc = self.make_formfunc(eq, F_target, arrays, targets) formrhs = self.make_rhs(eq, b, arrays) return (formfunc, formrhs) - def build_matvec_eqns(self, eq, target, arrays): + def build_matvec_eq(self, eq, target, arrays): b, F_target, _, targets = separate_eqn(eq, target) - if not F_target: - return None - matvec = self.make_matvec(eq, F_target, arrays, targets) - return matvec + if F_target: + return self.make_matvec(eq, F_target, arrays, targets) + return (None,) def make_matvec(self, eq, F_target, arrays, targets): if isinstance(eq, EssentialBC): @@ -103,39 +129,41 @@ def make_matvec(self, eq, F_target, arrays, targets): # Trivial equations for bc rows -> place 1.0 on diagonal (scaled) # and zero symmetrically. rhs = arrays['x'] - zero_column = Eq(arrays['x'], 0.0, subdomain=eq.subdomain) - return (EssentialBC(arrays['y'], rhs, subdomain=eq.subdomain), zero_column) + zero_row = ZeroRow(arrays['y'], rhs, subdomain=eq.subdomain) + zero_column = ZeroColumn(arrays['x'], 0.0, subdomain=eq.subdomain) + return (zero_row, zero_column) else: rhs = F_target.subs(targets_to_arrays(arrays['x'], targets)) - rhs = rhs.subs(self.time_mapper) - rhs = rhs * self.cell_area - return Eq(arrays['y'], rhs, subdomain=eq.subdomain) + rhs = rhs.subs(self.time_mapper) * self.cell_area + # TODO: Average centre stencils if they vary, to scale essential BC rows. + # self.centre = centre_stencil(rhs, arrays['x'], as_coeff=True) + # stencil = centre_stencil(rhs, arrays['x'], as_coeff=True) + # self._centre_stencils[arrays['x']].add(stencil) + # self._centre_stencils.add(stencil) + + return as_tuple(Eq(arrays['y'], rhs, subdomain=eq.subdomain)) def make_formfunc(self, eq, F_target, arrays, targets): if isinstance(eq, EssentialBC): - # The initial guess already satisfies the essential boundary conditions, - # so this term will always be zero. It's included here to allow - # testing of the Jacobian via finite differences. + # The initial guess satisfies the essential BCs, so this term is zero. + # Still included to support Jacobian testing via finite differences. rhs = arrays['x'] - eq.rhs - # Create equation to handle essential boundary conditions - we - # move the essential BCs to the right-hand side - # and zero the corresponding column in the Jacobian. - # TODO: extend this to mixed problems - move_bc = Eq(arrays['x'], eq.rhs, subdomain=eq.subdomain) - return (EssentialBC(arrays['f'], rhs, subdomain=eq.subdomain), move_bc) + zero_row = ZeroRow(arrays['f'], rhs, subdomain=eq.subdomain) + # Move essential boundary condition to the right-hand side + zero_col = ZeroColumn(arrays['x'], eq.rhs, subdomain=eq.subdomain) + return (zero_row, zero_col) else: if isinstance(F_target, (int, float)): rhs = F_target * self.cell_area else: rhs = F_target.subs(targets_to_arrays(arrays['x'], targets)) - rhs = rhs.subs(self.time_mapper) - rhs = rhs * self.cell_area - return Eq(arrays['f'], rhs, subdomain=eq.subdomain) + rhs = rhs.subs(self.time_mapper) * self.cell_area + return as_tuple(Eq(arrays['f'], rhs, subdomain=eq.subdomain)) def make_rhs(self, eq, b, arrays): rhs = 0. if isinstance(eq, EssentialBC) else b.subs(self.time_mapper) - rhs = rhs*self.cell_area - return Eq(arrays['b'], rhs, subdomain=eq.subdomain) + rhs = rhs * self.cell_area + return as_tuple(Eq(arrays['b'], rhs, subdomain=eq.subdomain)) def make_initial_guess(self, eq, target, arrays): """ @@ -162,9 +190,28 @@ def generate_arrays(self, target): for p in prefixes } + def scale_essential_bcs(self, equations, scale): + """ + Scale the essential boundary rows so that the Jacobian has a constant diagonal, + thereby reducing its condition number. + """ + # # stencils = self.centre_stencils[arrays['x']] + # if len(stencils) > 1: + # # Scaling of jacobian is therefore ambiguous, potentially could averge across the subblock + # # for now just set to trivial 1.0 + # scale = 1.0 + # else: + # scale = next(iter(stencils)) + return [ + eq._rebuild(rhs=scale * eq.rhs) if isinstance(eq, ZeroRow) else eq + for eq in equations + ] + class InjectSolveNested(InjectSolve): + def linear_solve_args(self): + combined_eqns = [] for eqns in self.target_eqns.values(): combined_eqns.extend(eqns) @@ -181,51 +228,90 @@ def linear_solve_args(self): self.cell_area = np.prod(all_data.grid.spacing_symbols) + all_formfuncs = [] + for target, eqns in self.target_eqns.items(): - eqns = as_tuple(eqns) - self.update_jacobian(eqns, target, jacobian, arrays[target]) - formfuncs = [ - self.build_function_eqns(eq, target, coupled_targets, arrays) - for eq in eqns - ] - all_data.extend_formfuncs(formfuncs) + # Update all rows of the Jacobian for this target + self.update_jacobian(as_tuple(eqns), target, jacobian, arrays[target]) + + formfuncs = chain.from_iterable( + self.build_function_eq(eq, target, coupled_targets, arrays) + for eq in as_tuple(eqns) + ) + # from IPython import embed; embed() + scale, = self._diag_scale[arrays[target]['x']] + all_formfuncs.extend(self.scale_essential_bcs(formfuncs, scale)) + + formfuncs = tuple(sorted( + all_formfuncs, key=lambda e: not isinstance(e, EssentialBC) + )) + all_data.extend_formfuncs(formfuncs) return target, tuple(funcs), all_data def update_jacobian(self, eqns, target, jacobian, arrays): + for submat, mtvs in jacobian.submatrices[target].items(): - # TODO: maintain symmetry for coupled - matvecs = [e for i in [self.build_matvec_eqns(eq, mtvs['derivative_wrt'], arrays) - for eq in eqns] for e in (i if isinstance(i, (tuple)) else [i])] + matvecs = [ + e for eq in eqns for e in + self.build_matvec_eq(eq, mtvs['derivative_wrt'], arrays) + ] matvecs = [m for m in matvecs if m is not None] + if submat in jacobian.diagonal_submatrix_keys: + stencils = set() + for eq in matvecs: + if not isinstance(eq, EssentialBC): + stencil = centre_stencil(eq.rhs, arrays['x'], as_coeff=True) + stencils.add(stencil) + # from IPython import embed; embed() + if len(stencils) > 1: + # Scaling of jacobian is therefore ambiguous, potentially could average across the subblock + # for now just set to trivial 1.0 + # TODO: doens't need to be a defaultdict, just a dict? + self._diag_scale[arrays['x']].add(1.0) + scale = 1.0 + else: + scale = next(iter(stencils)) + self._diag_scale[arrays['x']].add(scale) + # scale = next(iter(stencils)) + + matvecs = self.scale_essential_bcs(matvecs, scale) + + matvecs = tuple(sorted(matvecs, key=lambda e: not isinstance(e, EssentialBC))) + if matvecs: jacobian.set_submatrix(target, submat, matvecs) - def build_function_eqns(self, eq, main_target, coupled_targets, arrays): + def build_function_eq(self, eq, main_target, coupled_targets, arrays): zeroed = eq.lhs - eq.rhs - # TODO: clean up, test coupled with time dependence - zeroed_eqn = Eq(zeroed, 0) - zeroed_eqn = eval_time_derivatives(zeroed) + zeroed_eqn = Eq(eq.lhs - eq.rhs, 0) + eval_zeroed_eqn = eval_time_derivatives(zeroed_eqn.lhs) mapper = {} for t in coupled_targets: - target_funcs = generate_targets(Eq(zeroed, 0), t) + target_funcs = set(generate_targets(Eq(eval_zeroed_eqn, 0), t)) mapper.update(targets_to_arrays(arrays[t]['x'], target_funcs)) if isinstance(eq, EssentialBC): rhs = arrays[main_target]['x'] - eq.rhs + zero_row = ZeroRow( + arrays[main_target]['f'], rhs, subdomain=eq.subdomain + ) + zero_col = ZeroColumn( + arrays[main_target]['x'], eq.rhs, subdomain=eq.subdomain + ) + return (zero_row, zero_col) else: if isinstance(zeroed, (int, float)): rhs = zeroed * self.cell_area else: rhs = zeroed.subs(mapper) - rhs = rhs.subs(self.time_mapper) - rhs = rhs * self.cell_area + rhs = rhs.subs(self.time_mapper)*self.cell_area - return Eq(arrays[main_target]['f'], rhs, subdomain=eq.subdomain) + return as_tuple(Eq(arrays[main_target]['f'], rhs, subdomain=eq.subdomain)) def generate_arrays_combined(self, *targets): return { @@ -241,32 +327,37 @@ def generate_arrays_combined(self, *targets): for target in targets } - def make_matvec(self, eq, F_target, arrays, targets): - if isinstance(eq, EssentialBC): - # TODO: SHOULDN'T NEED this subclass once fixed mixed problems + symmetry - # NOTE: Until PetscSection + DMDA is supported, we leave - # the essential BCs in the solver. - # Trivial equations for bc rows -> place 1.0 on diagonal - # and zero symmetrically. FIX FOR MIXED PROBLEMS - rhs = arrays['x'] - else: - rhs = F_target.subs(targets_to_arrays(arrays['x'], targets)) - rhs = rhs.subs(self.time_mapper) - rhs = rhs * self.cell_area - return Eq(arrays['y'], rhs, subdomain=eq.subdomain) - class EssentialBC(Eq): """ - A special equation representing an essential boundary condition. - It is used to handle essential boundary conditions in the PETSc solver. - Until PetscSection + DMDA is supported, we treat essential boundary conditions - as trivial equations in the solver, where we place 1.0 on the diagonal of the jacobian, - zero symmetrically and move the essential boundary condition to the right-hand side. + A special equation used to handle essential boundary conditions + in the PETSc solver. Until PetscSection + DMDA is supported, + we treat essential boundary conditions as trivial equations + in the solver, where we place 1.0 (scaled) on the diagonal of + the jacobian, zero symmetrically and move the boundary + data to the right-hand side. NOTE: When users define essential boundary conditions, they need to ensure that the SubDomains do not overlap. Solver will still run but may see unexpected behaviour - at boundaries. + at boundaries. This will be documented in the PETSc examples. + """ + pass + + +class ZeroRow(EssentialBC): + """ + Equation used to zero the row of the Jacobian corresponding + to an essential BC. + This is only used interally by the compiler, not by users. + """ + pass + + +class ZeroColumn(EssentialBC): + """ + Equation used to zero the column of the Jacobian corresponding + to an essential BC. + This is only used interally by the compiler, not by users. """ pass @@ -354,26 +445,34 @@ def _(expr, targets): @singledispatch -def centre_stencil(expr, target): +def centre_stencil(expr, target, as_coeff=False): """ 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. + Parameters + ---------- + expr : the expression to extract the centre stencil from + target : the target function whose centre stencil we want + as_coeff : bool, optional + If True, return the coefficient of the centre stencil """ - return expr if expr == target else 0 + if expr == target: + return 1 if as_coeff else expr + return 0 @centre_stencil.register(sympy.Add) -def _(expr, target): +def _(expr, target, as_coeff=False): if not expr.has(target): return 0 - args = [centre_stencil(a, target) for a in expr.args] + args = [centre_stencil(a, target, as_coeff) for a in expr.args] return expr.func(*args, evaluate=False) @centre_stencil.register(Mul) -def _(expr, target): +def _(expr, target, as_coeff=False): if not expr.has(target): return 0 @@ -382,16 +481,16 @@ def _(expr, target): if not a.has(target): args.append(a) else: - args.append(centre_stencil(a, target)) + args.append(centre_stencil(a, target, as_coeff)) return expr.func(*args, evaluate=False) @centre_stencil.register(Derivative) -def _(expr, target): +def _(expr, target, as_coeff=False): if not expr.has(target): return 0 - args = [centre_stencil(a, target) for a in expr.evaluate.args] + args = [centre_stencil(a, target, as_coeff) for a in expr.evaluate.args] return expr.evaluate.func(*args) diff --git a/devito/petsc/types/array.py b/devito/petsc/types/array.py index 58ca3e28d7..049785662b 100644 --- a/devito/petsc/types/array.py +++ b/devito/petsc/types/array.py @@ -124,11 +124,16 @@ class PetscBundle(Bundle): is_Bundle = True _data_alignment = False + __rkwargs__ = Bundle.__rkwargs__ + ('pname',) + + def __init__(self, *args, pname="Field", **kwargs): + super().__init__(*args, **kwargs) + self._pname = pname + @property def _C_ctype(self): - # TODO: extend to cases with multiple petsc solves...(need diff struct name for each solve) fields = [(i.target.name, dtype_to_ctype(i.dtype)) for i in self.components] - return POINTER(type('Field', (Structure,), {'_fields_': fields})) + return POINTER(type(self.pname, (Structure,), {'_fields_': fields})) @cached_property def indexed(self): @@ -164,6 +169,10 @@ def __getitem__(self, index): raise ValueError("Expected %d or %d indices, got %d instead" % (self.ndim, self.ndim + 1, len(index))) + @property + def pname(self): + return self._pname + class AoSIndexedData(IndexedData): @property diff --git a/devito/petsc/types/types.py b/devito/petsc/types/types.py index afab9fff63..d2c45ca010 100644 --- a/devito/petsc/types/types.py +++ b/devito/petsc/types/types.py @@ -321,6 +321,18 @@ def get_submatrix(self, field, key): Retrieve a specific submatrix. """ return self.submatrices.get(field, {}).get(key, None) + + @property + def diagonal_submatrix_keys(self): + """ + Return a list of diagonal submatrix keys (e.g., ['J00', 'J11']). + """ + keys = [] + for i, target in enumerate(self.targets): + diag_key = f'J{i}{i}' + if diag_key in self.submatrices[target]: + keys.append(diag_key) + return keys def __repr__(self): return str(self.submatrices) diff --git a/devito/tools/utils.py b/devito/tools/utils.py index 0920298356..1285ffabec 100644 --- a/devito/tools/utils.py +++ b/devito/tools/utils.py @@ -346,4 +346,3 @@ def key(i): return (v, str(type(i))) return sorted(items, key=key, reverse=True) - diff --git a/examples/petsc/Poisson/01_poisson.py b/examples/petsc/Poisson/01_poisson.py index 6f24f74f64..70f17d94c6 100644 --- a/examples/petsc/Poisson/01_poisson.py +++ b/examples/petsc/Poisson/01_poisson.py @@ -66,7 +66,7 @@ def analytical(x, y): Ly = np.float64(1.) n_values = list(range(13, 174, 10)) -n_values = [9] +n_values = [7] dx = np.array([Lx/(n-1) for n in n_values]) errors = [] @@ -112,7 +112,7 @@ def analytical(x, y): error = np.linalg.norm(diff) / np.linalg.norm(phi_analytical[1:-1, 1:-1]) errors.append(error) -slope, _ = np.polyfit(np.log(dx), np.log(errors), 1) +# slope, _ = np.polyfit(np.log(dx), np.log(errors), 1) -assert slope > 1.9 -assert slope < 2.1 +# assert slope > 1.9 +# assert slope < 2.1 diff --git a/examples/petsc/random/biharmonic/02_biharmonic.py b/examples/petsc/random/biharmonic/02_biharmonic.py index bc7b54654d..973b16499e 100644 --- a/examples/petsc/random/biharmonic/02_biharmonic.py +++ b/examples/petsc/random/biharmonic/02_biharmonic.py @@ -12,9 +12,12 @@ configuration['compiler'] = 'custom' os.environ['CC'] = 'mpicc' + PetscInitialize() + # Subdomains to implement BCs + class SubTop(SubDomain): name = 'subtop' @@ -50,18 +53,23 @@ def define(self, dimensions): def c(x): return x**3 * (1 - x)**3 + def ddc(x): return 6.0 * x * (1 - x) * (1 - 5.0 * x + 5.0 * x**2) + def d4c(x): return -72.0 * (1 - 5.0 * x + 5.0 * x**2) + def u_exact_fcn(x, y): return c(x) * c(y) + def lap_u_exact_fcn(x, y): return -ddc(x) * c(y) - c(x) * ddc(y) + def f_fcn(x, y): return d4c(x) * c(y) + 2.0 * ddc(x) * ddc(y) + c(x) * d4c(y) @@ -77,6 +85,7 @@ def f_fcn(x, y): Ly = np.float64(1.) n_values = [33, 53, 73, 93, 113] +n_values = [9] dx = np.array([Lx/(n-1) for n in n_values]) u_errors = [] @@ -115,7 +124,8 @@ def f_fcn(x, y): bc_v += [EssentialBC(v, 0., subdomain=sub4)] # T (see ref) is nonsymmetric so need to set default KSP type to GMRES - petsc = PETScSolve({v: [eqn1]+bc_v, u: [eqn2]+bc_u}, solver_parameters={'ksp_rtol': 1e-10}) + params = {'ksp_rtol': 1e-10} + petsc = PETScSolve({v: [eqn1]+bc_v, u: [eqn2]+bc_u}, solver_parameters=params) with switchconfig(language='petsc'): op = Operator(petsc) @@ -126,12 +136,14 @@ def f_fcn(x, y): # Compute infinity norm for u u_diff = u_exact.data[:] - u.data[:] - u_error = np.linalg.norm(u_diff.ravel(), ord=np.inf) / np.linalg.norm(u_exact.data[:].ravel(), ord=np.inf) + u_diff_norm = np.linalg.norm(u_diff.ravel(), ord=np.inf) + u_error = u_diff_norm / np.linalg.norm(u_exact.data[:].ravel(), ord=np.inf) u_errors.append(u_error) # Compute infinity norm for lap_u v_diff = lap_u.data[:] - v.data[:] - v_error = np.linalg.norm(v_diff.ravel(), ord=np.inf) / np.linalg.norm(lap_u.data[:].ravel(), ord=np.inf) + v_diff_norm = np.linalg.norm(v_diff.ravel(), ord=np.inf) + v_error = v_diff_norm / np.linalg.norm(lap_u.data[:].ravel(), ord=np.inf) v_errors.append(v_error) u_slope, _ = np.polyfit(np.log(dx), np.log(u_errors), 1) diff --git a/tests/test_petsc.py b/tests/test_petsc.py index c136160ecc..382b325e4e 100644 --- a/tests/test_petsc.py +++ b/tests/test_petsc.py @@ -339,7 +339,7 @@ def test_separate_eqn(eqn, target, expected): f2 = TimeFunction(name='f2', grid=grid, space_order=so) # noqa g2 = TimeFunction(name='g2', grid=grid, space_order=so) # noqa - b, F, _, _= separate_eqn(eval(eqn), eval(target)) + b, F, _, _ = separate_eqn(eval(eqn), eval(target)) expected_b, expected_F = expected assert str(b) == expected_b @@ -793,8 +793,7 @@ def test_essential_bcs(): Verify that PETScSolve returns the correct output with essential boundary conditions. """ - - # SubDomains used for essential boundary conditions + # SubDomains used for essential boundary conditions # should not overlap. class SubTop(SubDomain): name = 'subtop' @@ -917,7 +916,6 @@ def test_coupled_vs_non_coupled(self): # Solving for multiple fields within the same matrix system requires # less callback functions than solving them separately. - # TODO: check reuse of callback functions where appropriate assert len(callbacks1) == 8 assert len(callbacks2) == 6 @@ -972,6 +970,10 @@ def test_coupled_structs(self): assert 'struct UserCtx0\n{' not in ccode assert 'struct UserCtx0\n{' in hcode + # The public struct Field0 only appears in the header file + assert 'struct Field0\n{' not in ccode + assert 'struct Field0\n{' in hcode + @skipif('petsc') def test_coupled_frees(self): grid = Grid(shape=(11, 11), dtype=np.float64) @@ -1083,7 +1085,7 @@ def test_submatrices(self): j00 = submatrices.get_submatrix(e, 'J00') j11 = submatrices.get_submatrix(g, 'J11') - # Compatible scaling to reduce condition number of jacobian + # Compatible scaling to reduce condition number of jacobian assert str(j00['matvecs'][0]) == 'Eq(y_e(x, y),' \ + ' h_x*h_y*(Derivative(x_e(x, y), (x, 2)) + Derivative(x_e(x, y), (y, 2))))' @@ -1096,13 +1098,98 @@ def test_submatrices(self): assert j10['derivative_wrt'] == e assert j11['derivative_wrt'] == g - # TODO: - # @skipif('petsc') - # def test_create_submats(self): + @skipif('petsc') + def test_residual_bundle(self): + grid = Grid(shape=(11, 11), dtype=np.float64) - # add tests for all new callbacks - # def test_create_whole_matmult(): + functions = [Function(name=n, grid=grid, space_order=2) + for n in ['e', 'f', 'g', 'h']] + e, f, g, h = functions + eq1 = Eq(e.laplace, h) + eq2 = Eq(f.laplace, h) + eq3 = Eq(g.laplace, h) -# // add coupled test with 3 targets -# // add coupled test with 1 target? + petsc1 = PETScSolve({e: [eq1]}) + petsc2 = PETScSolve({e: [eq1], f: [eq2]}) + petsc3 = PETScSolve({e: [eq1], f: [eq2], g: [eq3]}) + + with switchconfig(language='petsc'): + op1 = Operator(petsc1, opt='noop', name='op1') + op2 = Operator(petsc2, opt='noop', name='op2') + op3 = Operator(petsc3, opt='noop', name='op3') + + # Check pointers to array of Field structs. Note this is only + # required when dof>1 when constructing the multi-component DMDA. + f_aos = 'struct Field0 (* f_bundle)[info.gxm] = ' \ + + '(struct Field0 (*)[info.gxm]) f_bundle_vec;' + x_aos = 'struct Field0 (* x_bundle)[info.gxm] = ' \ + + '(struct Field0 (*)[info.gxm]) x_bundle_vec;' + + for op in (op1, op2, op3): + ccode = str(op.ccode) + assert f_aos in ccode + assert x_aos in ccode + + assert 'struct Field0\n{\n PetscScalar e;\n}' \ + in str(op1.ccode) + assert 'struct Field0\n{\n PetscScalar e;\n PetscScalar f;\n}' \ + in str(op2.ccode) + assert 'struct Field0\n{\n PetscScalar e;\n PetscScalar f;\n ' \ + + 'PetscScalar g;\n}' in str(op3.ccode) + + @skipif('petsc') + def test_essential_bcs(self): + """ + Test mixed problem with SubDomains + """ + class SubTop(SubDomain): + name = 'subtop' + + def define(self, dimensions): + x, y = dimensions + return {x: ('middle', 1, 1), y: ('right', 1)} + + sub1 = SubTop() + + grid = Grid(shape=(9, 9), subdomains=(sub1,), dtype=np.float64) + + u = Function(name='u', grid=grid, space_order=2) + v = Function(name='v', grid=grid, space_order=2) + f = Function(name='f', grid=grid, space_order=2) + + eqn1 = Eq(-v.laplace, f, subdomain=grid.interior) + eqn2 = Eq(-u.laplace, v, subdomain=grid.interior) + + bc_u = [EssentialBC(u, 0., subdomain=sub1)] + bc_v = [EssentialBC(v, 0., subdomain=sub1)] + + petsc = PETScSolve({v: [eqn1]+bc_v, u: [eqn2]+bc_u}) + + with switchconfig(language='petsc'): + op = Operator(petsc) + + # Test scaling + J00 = op._func_table['J00_MatMult0'].root + + # Essential BC row + assert 'a1[ix + 2][iy + 2] = (2.0/((o0->h_y*o0->h_y))' \ + ' + 2.0/((o0->h_x*o0->h_x)))*o0->h_x*o0->h_y*a0[ix + 2][iy + 2];' in str(J00) + # Check zeroing of essential BC columns + assert 'a0[ix + 2][iy + 2] = 0.0;' in str(J00) + # Interior loop + assert 'a1[ix + 2][iy + 2] = (2.0*(r0*a0[ix + 2][iy + 2] ' \ + '+ r1*a0[ix + 2][iy + 2]) - (r0*a0[ix + 1][iy + 2] + ' \ + 'r0*a0[ix + 3][iy + 2] + r1*a0[ix + 2][iy + 1] ' \ + '+ r1*a0[ix + 2][iy + 3]))*o0->h_x*o0->h_y;' in str(J00) + + # J00 and J11 are semantically identical so check efunc reuse + assert len(op._func_table.values()) == 7 + # J00_MatMult0 is reused (in replace of J11_MatMult0) + create = op._func_table['MatCreateSubMatrices0'].root + assert 'MatShellSetOperation(submat_arr[0],' \ + + 'MATOP_MULT,(void (*)(void))J00_MatMult0)' in str(create) + assert 'MatShellSetOperation(submat_arr[3],' \ + + 'MATOP_MULT,(void (*)(void))J00_MatMult0)' in str(create) + + # TODO: Test mixed, time dependent solvers From f7e00bbc1c1691665d14161a2dbddb947ce50ced Mon Sep 17 00:00:00 2001 From: ZoeLeibowitz Date: Tue, 3 Jun 2025 21:43:43 +0100 Subject: [PATCH 32/84] types: Edit PetscMixin and use it in other petsc classes --- devito/petsc/iet/routines.py | 18 ++++----- devito/petsc/solve.py | 6 +-- devito/petsc/types/object.py | 72 ++++++++++++++++-------------------- devito/petsc/types/types.py | 10 ++--- tests/test_petsc.py | 24 ++++++------ 5 files changed, 61 insertions(+), 69 deletions(-) diff --git a/devito/petsc/iet/routines.py b/devito/petsc/iet/routines.py index 7c824ad638..00c46f8b7d 100644 --- a/devito/petsc/iet/routines.py +++ b/devito/petsc/iet/routines.py @@ -653,8 +653,8 @@ def submatrices_callback(self): return self._submatrices_callback @property - def submatrices(self): - return self.injectsolve.expr.rhs.fielddata.submatrices + def jacobian(self): + return self.injectsolve.expr.rhs.fielddata.jacobian @property def main_matvec_callback(self): @@ -682,7 +682,7 @@ def _make_core(self): all_fielddata = injectsolve.expr.rhs.fielddata for t in targets: - row_matvecs = all_fielddata.submatrices.submatrices[t] + row_matvecs = all_fielddata.jacobian.submatrices[t] arrays = all_fielddata.arrays[t] for submat, mtvs in row_matvecs.items(): if mtvs['matvecs']: @@ -714,7 +714,7 @@ def _whole_matvec_body(self): jctx = objs['ljacctx'] ctx_main = petsc_call('MatShellGetContext', [objs['J'], Byref(jctx)]) - nonzero_submats = self.submatrices.nonzero_submatrix_keys + nonzero_submats = self.jacobian.nonzero_submatrix_keys zero_y_memory = petsc_call( 'VecSet', [objs['Y'], 0.0] @@ -722,7 +722,7 @@ def _whole_matvec_body(self): calls = () for sm in nonzero_submats: - idx = self.submatrices.submat_to_index[sm] + idx = self.jacobian.submat_to_index[sm] ctx = sobjs[f'{sm}ctx'] X = sobjs[f'{sm}X'] Y = sobjs[f'{sm}Y'] @@ -1012,14 +1012,14 @@ def _submat_callback_body(self): upper_bound = objs['nsubmats'] - 1 iteration = Iteration(List(body=iter_body), i, upper_bound) - nonzero_submats = self.submatrices.nonzero_submatrix_keys + nonzero_submats = self.jacobian.nonzero_submatrix_keys matvec_lookup = {mv.name.split('_')[0]: mv for mv in self.matvecs} matmult_op = [ petsc_call( 'MatShellSetOperation', [ - objs['submat_arr'].indexed[self.submatrices.submat_to_index[sb]], + objs['submat_arr'].indexed[self.jacobian.submat_to_index[sb]], 'MATOP_MULT', MatShellSetOp(matvec_lookup[sb].name, void, void), ], @@ -1157,8 +1157,8 @@ def _extend_build(self, base_dict): dim_labels[i]: PetscInt(dim_labels[i]) for i in range(space_dims) }) - submatrices = injectsolve.expr.rhs.fielddata.submatrices - submatrix_keys = submatrices.submatrix_keys + jacobian = injectsolve.expr.rhs.fielddata.jacobian + submatrix_keys = jacobian.submatrix_keys base_dict['jacctx'] = JacobianStruct( name=sreg.make_name(prefix=objs['ljacctx'].name), diff --git a/devito/petsc/solve.py b/devito/petsc/solve.py index 278a5092b8..218c146ac5 100644 --- a/devito/petsc/solve.py +++ b/devito/petsc/solve.py @@ -13,7 +13,7 @@ from devito.symbolics import retrieve_functions from devito.tools import as_tuple, filter_ordered from devito.petsc.types import (LinearSolveExpr, PETScArray, DMDALocalInfo, - FieldData, MultipleFieldData, SubMatrices) + FieldData, MultipleFieldData, Jacobian) __all__ = ['PETScSolve', 'EssentialBC'] @@ -219,11 +219,11 @@ def linear_solve_args(self): self.time_mapper = generate_time_mapper(funcs) coupled_targets = list(self.target_eqns.keys()) - jacobian = SubMatrices(coupled_targets) + jacobian = Jacobian(coupled_targets) arrays = self.generate_arrays_combined(*coupled_targets) - all_data = MultipleFieldData(submatrices=jacobian, arrays=arrays, + all_data = MultipleFieldData(jacobian=jacobian, arrays=arrays, targets=coupled_targets) self.cell_area = np.prod(all_data.grid.spacing_symbols) diff --git a/devito/petsc/types/object.py b/devito/petsc/types/object.py index a285863d14..e62d125dc2 100644 --- a/devito/petsc/types/object.py +++ b/devito/petsc/types/object.py @@ -10,10 +10,17 @@ class PetscMixin: @property def _C_free_priority(self): - return FREE_PRIORITY[self] + if type(self) in FREE_PRIORITY: + return FREE_PRIORITY[type(self)] + else: + return super()._C_free_priority -class CallbackDM(LocalObject): +class PetscObject(PetscMixin, LocalObject): + pass + + +class CallbackDM(PetscObject): """ PETSc Data Management object (DM). This is the DM instance accessed within the callback functions via `SNESGetDM` and @@ -22,14 +29,12 @@ class CallbackDM(LocalObject): dtype = CustomDtype('DM') -class DM(LocalObject, PetscMixin): +class DM(CallbackDM): """ PETSc Data Management object (DM). This is the primary DM instance created within the main kernel and linked to the SNES solver using `SNESSetDM`. """ - dtype = CustomDtype('DM') - def __init__(self, *args, dofs=1, **kwargs): super().__init__(*args, **kwargs) self._dofs = dofs @@ -46,7 +51,7 @@ def _C_free(self): DMCast = cast('DM') -class CallbackMat(LocalObject): +class CallbackMat(PetscObject): """ PETSc Matrix object (Mat) used within callback functions. These instances are not destroyed during callback execution; @@ -55,19 +60,13 @@ class CallbackMat(LocalObject): dtype = CustomDtype('Mat') -class Mat(LocalObject): - dtype = CustomDtype('Mat') - +class Mat(CallbackMat): @property def _C_free(self): return petsc_call('MatDestroy', [Byref(self.function)]) - @property - def _C_free_priority(self): - return 2 - -class CallbackVec(LocalObject): +class CallbackVec(PetscObject): """ PETSc vector object (Vec). """ @@ -79,12 +78,8 @@ class Vec(CallbackVec): def _C_free(self): return petsc_call('VecDestroy', [Byref(self.function)]) - @property - def _C_free_priority(self): - return 1 - -class PetscMPIInt(LocalObject): +class PetscMPIInt(PetscObject): """ PETSc datatype used to represent `int` parameters to MPI functions. @@ -92,7 +87,7 @@ class PetscMPIInt(LocalObject): dtype = CustomDtype('PetscMPIInt') -class PetscInt(LocalObject): +class PetscInt(PetscObject): """ PETSc datatype used to represent `int` parameters to PETSc functions. @@ -100,7 +95,7 @@ class PetscInt(LocalObject): dtype = CustomIntType('PetscInt') -class KSP(LocalObject): +class KSP(PetscObject): """ PETSc KSP : Linear Systems Solvers. Manages Krylov Methods. @@ -108,7 +103,7 @@ class KSP(LocalObject): dtype = CustomDtype('KSP') -class CallbackSNES(LocalObject): +class CallbackSNES(PetscObject): """ PETSc SNES : Non-Linear Systems Solvers. """ @@ -120,19 +115,15 @@ class SNES(CallbackSNES): def _C_free(self): return petsc_call('SNESDestroy', [Byref(self.function)]) - @property - def _C_free_priority(self): - return 3 - -class PC(LocalObject): +class PC(PetscObject): """ PETSc object that manages all preconditioners (PC). """ dtype = CustomDtype('PC') -class KSPConvergedReason(LocalObject): +class KSPConvergedReason(PetscObject): """ PETSc object - reason a Krylov method was determined to have converged or diverged. @@ -140,7 +131,7 @@ class KSPConvergedReason(LocalObject): dtype = CustomDtype('KSPConvergedReason') -class DMDALocalInfo(LocalObject): +class DMDALocalInfo(PetscObject): """ PETSc object - C struct containing information about the local grid. @@ -148,7 +139,7 @@ class DMDALocalInfo(LocalObject): dtype = CustomDtype('DMDALocalInfo') -class PetscErrorCode(LocalObject): +class PetscErrorCode(PetscObject): """ PETSc datatype used to return PETSc error codes. https://petsc.org/release/manualpages/Sys/PetscErrorCode/ @@ -156,7 +147,7 @@ class PetscErrorCode(LocalObject): dtype = CustomDtype('PetscErrorCode') -class DummyArg(LocalObject): +class DummyArg(PetscObject): """ A void pointer used to satisfy the function signature of the `FormFunction` callback. @@ -164,21 +155,21 @@ class DummyArg(LocalObject): dtype = CustomDtype('void', modifier='*') -class MatReuse(LocalObject): +class MatReuse(PetscObject): dtype = CustomDtype('MatReuse') -class VecScatter(LocalObject): +class VecScatter(PetscObject): dtype = CustomDtype('VecScatter') -class StartPtr(LocalObject): +class StartPtr(PetscObject): def __init__(self, name, dtype): super().__init__(name=name) self.dtype = POINTER(dtype_to_ctype(dtype)) -class SingleIS(LocalObject): +class SingleIS(PetscObject): dtype = CustomDtype('IS') @@ -218,7 +209,8 @@ def __init__(self, name='subctx', pname='SubMatrixCtx', fields=None, _C_modifier = None -class PETScArrayObject(ArrayObject): + +class PETScArrayObject(PetscMixin, ArrayObject): _data_alignment = False def __init_finalize__(self, *args, **kwargs): @@ -251,10 +243,6 @@ def _C_name(self): def _mem_stack(self): return False - @property - def _C_free_priority(self): - return 0 - class CallbackPointerIS(PETScArrayObject): """ @@ -309,5 +297,9 @@ def _C_ctype(self): FREE_PRIORITY = { + PETScArrayObject: 0, + Vec: 1, + Mat: 2, + SNES: 3, DM: 4, } diff --git a/devito/petsc/types/types.py b/devito/petsc/types/types.py index d2c45ca010..75ee8e8254 100644 --- a/devito/petsc/types/types.py +++ b/devito/petsc/types/types.py @@ -189,10 +189,10 @@ def targets(self): class MultipleFieldData(FieldData): - def __init__(self, targets, arrays, submatrices=None): + def __init__(self, targets, arrays, jacobian=None): self._targets = as_tuple(targets) self._arrays = arrays - self._submatrices = submatrices + self._jacobian = jacobian self._formfuncs = [] def extend_formfuncs(self, formfuncs): @@ -232,8 +232,8 @@ def space_order(self): return space_orders.pop() @property - def submatrices(self): - return self._submatrices + def jacobian(self): + return self._jacobian @property def targets(self): @@ -244,7 +244,7 @@ def arrays(self): return self._arrays -class SubMatrices: +class Jacobian: def __init__(self, targets): self.targets = targets self.submatrices = self._initialize_submatrices() diff --git a/tests/test_petsc.py b/tests/test_petsc.py index 382b325e4e..8bcaeb3fa6 100644 --- a/tests/test_petsc.py +++ b/tests/test_petsc.py @@ -1065,25 +1065,25 @@ def test_submatrices(self): petsc = PETScSolve({e: [eq1], g: [eq2]}) - submatrices = petsc[0].rhs.fielddata.submatrices + jacobian = petsc[0].rhs.fielddata.jacobian - j00 = submatrices.get_submatrix(e, 'J00') - j01 = submatrices.get_submatrix(e, 'J01') - j10 = submatrices.get_submatrix(g, 'J10') - j11 = submatrices.get_submatrix(g, 'J11') + j00 = jacobian.get_submatrix(e, 'J00') + j01 = jacobian.get_submatrix(e, 'J01') + j10 = jacobian.get_submatrix(g, 'J10') + j11 = jacobian.get_submatrix(g, 'J11') # Check the number of submatrices - assert len(submatrices.submatrix_keys) == 4 - assert str(submatrices.submatrix_keys) == "['J00', 'J01', 'J10', 'J11']" + assert len(jacobian.submatrix_keys) == 4 + assert str(jacobian.submatrix_keys) == "['J00', 'J01', 'J10', 'J11']" # Technically a non-coupled problem, so the only non-zero submatrices # should be the diagonal ones i.e J00 and J11 - assert submatrices.nonzero_submatrix_keys == ['J00', 'J11'] - assert submatrices.get_submatrix(e, 'J01')['matvecs'] is None - assert submatrices.get_submatrix(g, 'J10')['matvecs'] is None + assert jacobian.nonzero_submatrix_keys == ['J00', 'J11'] + assert jacobian.get_submatrix(e, 'J01')['matvecs'] is None + assert jacobian.get_submatrix(g, 'J10')['matvecs'] is None - j00 = submatrices.get_submatrix(e, 'J00') - j11 = submatrices.get_submatrix(g, 'J11') + j00 = jacobian.get_submatrix(e, 'J00') + j11 = jacobian.get_submatrix(g, 'J11') # Compatible scaling to reduce condition number of jacobian assert str(j00['matvecs'][0]) == 'Eq(y_e(x, y),' \ From 6e20d4bad6fa26a8167ce1edf8a1186d859bd66c Mon Sep 17 00:00:00 2001 From: ZoeLeibowitz Date: Tue, 3 Jun 2025 22:32:56 +0100 Subject: [PATCH 33/84] dsl: Add extraction file --- devito/symbolics/extraction.py | 120 +++++++++++++++++++++++++++++++++ 1 file changed, 120 insertions(+) create mode 100644 devito/symbolics/extraction.py diff --git a/devito/symbolics/extraction.py b/devito/symbolics/extraction.py new file mode 100644 index 0000000000..dfbc145748 --- /dev/null +++ b/devito/symbolics/extraction.py @@ -0,0 +1,120 @@ +from functools import singledispatch + +import sympy + +from devito.finite_differences.differentiable import Mul +from devito.finite_differences.derivative import Derivative +from devito.types.equation import Eq +from devito.symbolics import retrieve_functions + + +__all__ = ['separate_eqn', 'generate_targets', 'centre_stencil'] + + +def separate_eqn(eqn, target): + """ + Separate the equation into two separate expressions, + where F(target) = b. + """ + zeroed_eqn = Eq(eqn.lhs - eqn.rhs, 0) + from devito.operations.solve import eval_time_derivatives + zeroed_eqn = eval_time_derivatives(zeroed_eqn.lhs) + target_funcs = set(generate_targets(zeroed_eqn, target)) + b, F_target = remove_targets(zeroed_eqn, target_funcs) + return -b, F_target, zeroed_eqn, target_funcs + + +def generate_targets(eq, target): + """ + Extract all the functions that share the same time index as the target + but may have different spatial indices. + """ + funcs = retrieve_functions(eq) + if target.is_TimeFunction: + time_idx = target.indices[target.time_dim] + targets = [ + f for f in funcs if f.function is target.function and time_idx + in f.indices + ] + else: + targets = [f for f in funcs if f.function is target.function] + return targets + + +@singledispatch +def remove_targets(expr, targets): + return (0, expr) if expr in targets else (expr, 0) + + +@remove_targets.register(sympy.Add) +def _(expr, targets): + if not any(expr.has(t) for t in targets): + return (expr, 0) + + args_b, args_F = zip(*(remove_targets(a, targets) for a in expr.args)) + return (expr.func(*args_b, evaluate=False), expr.func(*args_F, evaluate=False)) + + +@remove_targets.register(Mul) +def _(expr, targets): + if not any(expr.has(t) for t in targets): + return (expr, 0) + + args_b, args_F = zip(*[remove_targets(a, targets) if any(a.has(t) for t in targets) + else (a, a) for a in expr.args]) + return (expr.func(*args_b, evaluate=False), expr.func(*args_F, evaluate=False)) + + +@remove_targets.register(Derivative) +def _(expr, targets): + return (0, expr) if any(expr.has(t) for t in targets) else (expr, 0) + + +@singledispatch +def centre_stencil(expr, target, as_coeff=False): + """ + 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. + Parameters + ---------- + expr : the expression to extract the centre stencil from + target : the target function whose centre stencil we want + as_coeff : bool, optional + If True, return the coefficient of the centre stencil + """ + if expr == target: + return 1 if as_coeff else expr + return 0 + + +@centre_stencil.register(sympy.Add) +def _(expr, target, as_coeff=False): + if not expr.has(target): + return 0 + + args = [centre_stencil(a, target, as_coeff) for a in expr.args] + return expr.func(*args, evaluate=False) + + +@centre_stencil.register(Mul) +def _(expr, target, as_coeff=False): + 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, as_coeff)) + + return expr.func(*args, evaluate=False) + + +@centre_stencil.register(Derivative) +def _(expr, target, as_coeff=False): + if not expr.has(target): + return 0 + args = [centre_stencil(a, target, as_coeff) for a in expr.evaluate.args] + return expr.evaluate.func(*args) From 8a79fcde140004230b2bfe79323f1db7a17fc2b2 Mon Sep 17 00:00:00 2001 From: ZoeLeibowitz Date: Tue, 3 Jun 2025 22:33:35 +0100 Subject: [PATCH 34/84] dsl: Move symbolic extraction functions outside of the PETSc module --- devito/petsc/solve.py | 136 ++++--------------- devito/symbolics/__init__.py | 1 + tests/test_petsc.py | 251 +---------------------------------- tests/test_symbolics.py | 241 ++++++++++++++++++++++++++++++++- 4 files changed, 271 insertions(+), 358 deletions(-) diff --git a/devito/petsc/solve.py b/devito/petsc/solve.py index 218c146ac5..f3e7cd377b 100644 --- a/devito/petsc/solve.py +++ b/devito/petsc/solve.py @@ -11,6 +11,8 @@ from devito.types.equation import PetscEq from devito.operations.solve import eval_time_derivatives from devito.symbolics import retrieve_functions +from devito.symbolics.extraction import (separate_eqn, centre_stencil, + generate_targets) from devito.tools import as_tuple, filter_ordered from devito.petsc.types import (LinearSolveExpr, PETScArray, DMDALocalInfo, FieldData, MultipleFieldData, Jacobian) @@ -76,7 +78,6 @@ def generate_field_data(self, eqns, target, arrays): zip(*[self.build_function_eq(eq, target, arrays) for eq in eqns]) ) - # self._centre_stencils[arrays['x']].update( stencils = set() for eq in matvecs: if not isinstance(eq, EssentialBC): @@ -362,33 +363,33 @@ class ZeroColumn(EssentialBC): pass -def separate_eqn(eqn, target): - """ - Separate the equation into two separate expressions, - where F(target) = b. - """ - zeroed_eqn = Eq(eqn.lhs - eqn.rhs, 0) - zeroed_eqn = eval_time_derivatives(zeroed_eqn.lhs) - target_funcs = set(generate_targets(zeroed_eqn, target)) - b, F_target = remove_targets(zeroed_eqn, target_funcs) - return -b, F_target, zeroed_eqn, target_funcs - - -def generate_targets(eq, target): - """ - Extract all the functions that share the same time index as the target - but may have different spatial indices. - """ - funcs = retrieve_functions(eq) - if isinstance(target, TimeFunction): - time_idx = target.indices[target.time_dim] - targets = [ - f for f in funcs if f.function is target.function and time_idx - in f.indices - ] - else: - targets = [f for f in funcs if f.function is target.function] - return targets +# def separate_eqn(eqn, target): +# """ +# Separate the equation into two separate expressions, +# where F(target) = b. +# """ +# zeroed_eqn = Eq(eqn.lhs - eqn.rhs, 0) +# zeroed_eqn = eval_time_derivatives(zeroed_eqn.lhs) +# target_funcs = set(generate_targets(zeroed_eqn, target)) +# b, F_target = remove_targets(zeroed_eqn, target_funcs) +# return -b, F_target, zeroed_eqn, target_funcs + + +# def generate_targets(eq, target): +# """ +# Extract all the functions that share the same time index as the target +# but may have different spatial indices. +# """ +# funcs = retrieve_functions(eq) +# if isinstance(target, TimeFunction): +# time_idx = target.indices[target.time_dim] +# targets = [ +# f for f in funcs if f.function is target.function and time_idx +# in f.indices +# ] +# else: +# targets = [f for f in funcs if f.function is target.function] +# return targets def targets_to_arrays(array, targets): @@ -415,85 +416,6 @@ def targets_to_arrays(array, targets): return dict(zip(targets, array_targets)) -@singledispatch -def remove_targets(expr, targets): - return (0, expr) if expr in targets else (expr, 0) - - -@remove_targets.register(sympy.Add) -def _(expr, targets): - if not any(expr.has(t) for t in targets): - return (expr, 0) - - args_b, args_F = zip(*(remove_targets(a, targets) for a in expr.args)) - return (expr.func(*args_b, evaluate=False), expr.func(*args_F, evaluate=False)) - - -@remove_targets.register(Mul) -def _(expr, targets): - if not any(expr.has(t) for t in targets): - return (expr, 0) - - args_b, args_F = zip(*[remove_targets(a, targets) if any(a.has(t) for t in targets) - else (a, a) for a in expr.args]) - return (expr.func(*args_b, evaluate=False), expr.func(*args_F, evaluate=False)) - - -@remove_targets.register(Derivative) -def _(expr, targets): - return (0, expr) if any(expr.has(t) for t in targets) else (expr, 0) - - -@singledispatch -def centre_stencil(expr, target, as_coeff=False): - """ - 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. - Parameters - ---------- - expr : the expression to extract the centre stencil from - target : the target function whose centre stencil we want - as_coeff : bool, optional - If True, return the coefficient of the centre stencil - """ - if expr == target: - return 1 if as_coeff else expr - return 0 - - -@centre_stencil.register(sympy.Add) -def _(expr, target, as_coeff=False): - if not expr.has(target): - return 0 - - args = [centre_stencil(a, target, as_coeff) for a in expr.args] - return expr.func(*args, evaluate=False) - - -@centre_stencil.register(Mul) -def _(expr, target, as_coeff=False): - 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, as_coeff)) - - return expr.func(*args, evaluate=False) - - -@centre_stencil.register(Derivative) -def _(expr, target, as_coeff=False): - if not expr.has(target): - return 0 - args = [centre_stencil(a, target, as_coeff) 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 diff --git a/devito/symbolics/__init__.py b/devito/symbolics/__init__.py index 3f1525297a..47935f789f 100644 --- a/devito/symbolics/__init__.py +++ b/devito/symbolics/__init__.py @@ -4,3 +4,4 @@ from devito.symbolics.search import * # noqa from devito.symbolics.inspection import * # noqa from devito.symbolics.manipulation import * # noqa +from devito.symbolics.extraction import * # noqa diff --git a/tests/test_petsc.py b/tests/test_petsc.py index 8bcaeb3fa6..8f9cc4d9a8 100644 --- a/tests/test_petsc.py +++ b/tests/test_petsc.py @@ -12,8 +12,7 @@ from devito.petsc.types import (DM, Mat, Vec, PetscMPIInt, KSP, PC, KSPConvergedReason, PETScArray, LinearSolveExpr, FieldData, MultipleFieldData) -from devito.petsc.solve import (PETScSolve, separate_eqn, centre_stencil, - EssentialBC) +from devito.petsc.solve import PETScSolve, EssentialBC from devito.petsc.iet.nodes import Expression from devito.petsc.initialize import PetscInitialize @@ -276,254 +275,6 @@ def test_cinterface_petsc_struct(): assert 'struct UserCtx0\n{' in hcode -@skipif('petsc') -@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. - - 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)) - - 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') -@pytest.mark.parametrize('eqn, target, expected', [ - ('Eq(f1.laplace, g1).evaluate', 'f1', - ( - 'g1(x, y)', - '-2.0*f1(x, y)/h_x**2 + f1(x - h_x, y)/h_x**2 + f1(x + h_x, y)/h_x**2 ' - '- 2.0*f1(x, y)/h_y**2 + f1(x, y - h_y)/h_y**2 + f1(x, y + h_y)/h_y**2' - )), - ('Eq(g1, f1.laplace).evaluate', 'f1', - ( - '-g1(x, y)', - '-(-2.0*f1(x, y)/h_x**2 + f1(x - h_x, y)/h_x**2 + f1(x + h_x, y)/h_x**2) ' - '- (-2.0*f1(x, y)/h_y**2 + f1(x, y - h_y)/h_y**2 + f1(x, y + h_y)/h_y**2)' - )), - ('Eq(g1, f1.laplace).evaluate', 'g1', - ( - '-2.0*f1(x, y)/h_x**2 + f1(x - h_x, y)/h_x**2 + f1(x + h_x, y)/h_x**2 ' - '- 2.0*f1(x, y)/h_y**2 + f1(x, y - h_y)/h_y**2 + f1(x, y + h_y)/h_y**2', - 'g1(x, y)' - )), - ('Eq(f1 + f1.laplace, g1).evaluate', 'f1', - ( - 'g1(x, y)', - '-2.0*f1(x, y)/h_x**2 + f1(x - h_x, y)/h_x**2 + f1(x + h_x, y)/h_x**2 - 2.0' - '*f1(x, y)/h_y**2 + f1(x, y - h_y)/h_y**2 + f1(x, y + h_y)/h_y**2 + f1(x, y)' - )), - ('Eq(g1.dx + f1.dx, g1).evaluate', 'f1', - ( - '-(-g1(x, y)/h_x + g1(x + h_x, y)/h_x) + g1(x, y)', - '-f1(x, y)/h_x + f1(x + h_x, y)/h_x' - )), - ('Eq(g1.dx + f1.dx, g1).evaluate', 'g1', - ( - '-(-f1(x, y)/h_x + f1(x + h_x, y)/h_x)', - '-g1(x, y)/h_x + g1(x + h_x, y)/h_x - g1(x, y)' - )), - ('Eq(f1 * g1.dx, g1).evaluate', 'g1', - ( - '0', '(-g1(x, y)/h_x + g1(x + h_x, y)/h_x)*f1(x, y) - g1(x, y)' - )), - ('Eq(f1 * g1.dx, g1).evaluate', 'f1', - ( - 'g1(x, y)', '(-g1(x, y)/h_x + g1(x + h_x, y)/h_x)*f1(x, y)' - )), - ('Eq((f1 * g1.dx).dy, f1).evaluate', 'f1', - ( - '0', '(-1/h_y)*(-g1(x, y)/h_x + g1(x + h_x, y)/h_x)*f1(x, y) ' - '+ (-g1(x, y + h_y)/h_x + g1(x + h_x, y + h_y)/h_x)*f1(x, y + h_y)/h_y ' - '- f1(x, y)' - )), - ('Eq((f1 * g1.dx).dy, f1).evaluate', 'g1', - ( - 'f1(x, y)', '(-1/h_y)*(-g1(x, y)/h_x + g1(x + h_x, y)/h_x)*f1(x, y) + ' - '(-g1(x, y + h_y)/h_x + g1(x + h_x, y + h_y)/h_x)*f1(x, y + h_y)/h_y' - )), - ('Eq(f2.laplace, g2).evaluate', 'g2', - ( - '-(-2.0*f2(t, x, y)/h_x**2 + f2(t, x - h_x, y)/h_x**2 + f2(t, x + h_x, y)' - '/h_x**2) - (-2.0*f2(t, x, y)/h_y**2 + f2(t, x, y - h_y)/h_y**2 + ' - 'f2(t, x, y + h_y)/h_y**2)', '-g2(t, x, y)' - )), - ('Eq(f2.laplace, g2).evaluate', 'f2', - ( - 'g2(t, x, y)', '-2.0*f2(t, x, y)/h_x**2 + f2(t, x - h_x, y)/h_x**2 + ' - 'f2(t, x + h_x, y)/h_x**2 - 2.0*f2(t, x, y)/h_y**2 + f2(t, x, y - h_y)' - '/h_y**2 + f2(t, x, y + h_y)/h_y**2' - )), - ('Eq(f2.laplace, f2).evaluate', 'f2', - ( - '0', '-2.0*f2(t, x, y)/h_x**2 + f2(t, x - h_x, y)/h_x**2 + ' - 'f2(t, x + h_x, y)/h_x**2 - 2.0*f2(t, x, y)/h_y**2 + f2(t, x, y - h_y)/h_y**2' - ' + f2(t, x, y + h_y)/h_y**2 - f2(t, x, y)' - )), - ('Eq(g2*f2.laplace, f2).evaluate', 'g2', - ( - 'f2(t, x, y)', '(-2.0*f2(t, x, y)/h_x**2 + f2(t, x - h_x, y)/h_x**2 + ' - 'f2(t, x + h_x, y)/h_x**2 - 2.0*f2(t, x, y)/h_y**2 + f2(t, x, y - h_y)/h_y**2' - ' + f2(t, x, y + h_y)/h_y**2)*g2(t, x, y)' - )), - ('Eq(f2.forward.laplace, f2).evaluate', 'f2.forward', - ( - 'f2(t, x, y)', '-2.0*f2(t + dt, x, y)/h_x**2 + f2(t + dt, x - h_x, y)/h_x**2' - ' + f2(t + dt, x + h_x, y)/h_x**2 - 2.0*f2(t + dt, x, y)/h_y**2 + ' - 'f2(t + dt, x, y - h_y)/h_y**2 + f2(t + dt, x, y + h_y)/h_y**2' - )), - ('Eq(f2.forward.laplace, f2).evaluate', 'f2', - ( - '-(-2.0*f2(t + dt, x, y)/h_x**2 + f2(t + dt, x - h_x, y)/h_x**2 + ' - 'f2(t + dt, x + h_x, y)/h_x**2) - (-2.0*f2(t + dt, x, y)/h_y**2 + ' - 'f2(t + dt, x, y - h_y)/h_y**2 + f2(t + dt, x, y + h_y)/h_y**2)', - '-f2(t, x, y)' - )), - ('Eq(f2.laplace + f2.forward.laplace, g2).evaluate', 'f2.forward', - ( - '-(-2.0*f2(t, x, y)/h_x**2 + f2(t, x - h_x, y)/h_x**2 + f2(t, x + h_x, y)/' - 'h_x**2) - (-2.0*f2(t, x, y)/h_y**2 + f2(t, x, y - h_y)/h_y**2 + ' - 'f2(t, x, y + h_y)/h_y**2) + g2(t, x, y)', '-2.0*f2(t + dt, x, y)/h_x**2 + ' - 'f2(t + dt, x - h_x, y)/h_x**2 + f2(t + dt, x + h_x, y)/h_x**2 - 2.0*' - 'f2(t + dt, x, y)/h_y**2 + f2(t + dt, x, y - h_y)/h_y**2 + ' - 'f2(t + dt, x, y + h_y)/h_y**2' - )), - ('Eq(g2.laplace, f2 + g2.forward).evaluate', 'g2.forward', - ( - '-(-2.0*g2(t, x, y)/h_x**2 + g2(t, x - h_x, y)/h_x**2 + ' - 'g2(t, x + h_x, y)/h_x**2) - (-2.0*g2(t, x, y)/h_y**2 + g2(t, x, y - h_y)' - '/h_y**2 + g2(t, x, y + h_y)/h_y**2) + f2(t, x, y)', '-g2(t + dt, x, y)' - )) -]) -def test_separate_eval_eqn(eqn, target, expected): - """ - Test the separate_eqn function on pre-evaluated equations. - This ensures that evaluated equations can be passed to PETScSolve, - allowing users to modify stencils for specific boundary conditions, - such as implementing free surface boundary conditions. - """ - grid = Grid((2, 2)) - - so = 2 - - f1 = Function(name='f1', grid=grid, space_order=so) # noqa - g1 = Function(name='g1', grid=grid, space_order=so) # noqa - - f2 = TimeFunction(name='f2', grid=grid, space_order=so) # noqa - g2 = TimeFunction(name='g2', grid=grid, space_order=so) # noqa - - b, F, _, _ = separate_eqn(eval(eqn), eval(target)) - expected_b, expected_F = expected - - assert str(b) == expected_b - assert str(F) == expected_F - - -@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 - - f2 = TimeFunction(name='f2', grid=grid, space_order=so) # noqa - g2 = TimeFunction(name='g2', grid=grid, space_order=so) # noqa - - centre = centre_stencil(eval(expr), eval(target)) - - assert str(centre) == expected - - @skipif('petsc') def test_callback_arguments(): """ diff --git a/tests/test_symbolics.py b/tests/test_symbolics.py index deedbec950..1d55ce5a59 100644 --- a/tests/test_symbolics.py +++ b/tests/test_symbolics.py @@ -14,7 +14,8 @@ CallFromPointer, Cast, DefFunction, FieldFromPointer, INT, FieldFromComposite, IntDiv, Namespace, Rvalue, ReservedWord, ListInitializer, uxreplace, pow_to_mul, - retrieve_derivatives, BaseCast, SizeOf, sympy_dtype) + retrieve_derivatives, BaseCast, SizeOf, sympy_dtype, + separate_eqn, centre_stencil) from devito.tools import as_tuple from devito.types import (Array, Bundle, FIndexed, LocalObject, Object, ComponentAccess, StencilDimension, Symbol as dSymbol, @@ -927,3 +928,241 @@ def test_print_div(): b = SizeOf(np.int64) cstr = ccode(a / b) assert cstr == 'sizeof(int)/sizeof(long)' + + +@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. + """ + grid = Grid((2, 2)) + + so = 2 + + f1 = Function(name='f1', grid=grid, space_order=so) # noqa + g1 = Function(name='g1', grid=grid, space_order=so) # noqa + + f2 = TimeFunction(name='f2', grid=grid, space_order=so) # noqa + g2 = TimeFunction(name='g2', grid=grid, space_order=so) # noqa + + b, F, _, _ = separate_eqn(eval(eqn), eval(target)) + expected_b, expected_F = expected + + assert str(b) == expected_b + assert str(F) == expected_F + + +@pytest.mark.parametrize('eqn, target, expected', [ + ('Eq(f1.laplace, g1).evaluate', 'f1', + ( + 'g1(x, y)', + '-2.0*f1(x, y)/h_x**2 + f1(x - h_x, y)/h_x**2 + f1(x + h_x, y)/h_x**2 ' + '- 2.0*f1(x, y)/h_y**2 + f1(x, y - h_y)/h_y**2 + f1(x, y + h_y)/h_y**2' + )), + ('Eq(g1, f1.laplace).evaluate', 'f1', + ( + '-g1(x, y)', + '-(-2.0*f1(x, y)/h_x**2 + f1(x - h_x, y)/h_x**2 + f1(x + h_x, y)/h_x**2) ' + '- (-2.0*f1(x, y)/h_y**2 + f1(x, y - h_y)/h_y**2 + f1(x, y + h_y)/h_y**2)' + )), + ('Eq(g1, f1.laplace).evaluate', 'g1', + ( + '-2.0*f1(x, y)/h_x**2 + f1(x - h_x, y)/h_x**2 + f1(x + h_x, y)/h_x**2 ' + '- 2.0*f1(x, y)/h_y**2 + f1(x, y - h_y)/h_y**2 + f1(x, y + h_y)/h_y**2', + 'g1(x, y)' + )), + ('Eq(f1 + f1.laplace, g1).evaluate', 'f1', + ( + 'g1(x, y)', + '-2.0*f1(x, y)/h_x**2 + f1(x - h_x, y)/h_x**2 + f1(x + h_x, y)/h_x**2 - 2.0' + '*f1(x, y)/h_y**2 + f1(x, y - h_y)/h_y**2 + f1(x, y + h_y)/h_y**2 + f1(x, y)' + )), + ('Eq(g1.dx + f1.dx, g1).evaluate', 'f1', + ( + '-(-g1(x, y)/h_x + g1(x + h_x, y)/h_x) + g1(x, y)', + '-f1(x, y)/h_x + f1(x + h_x, y)/h_x' + )), + ('Eq(g1.dx + f1.dx, g1).evaluate', 'g1', + ( + '-(-f1(x, y)/h_x + f1(x + h_x, y)/h_x)', + '-g1(x, y)/h_x + g1(x + h_x, y)/h_x - g1(x, y)' + )), + ('Eq(f1 * g1.dx, g1).evaluate', 'g1', + ( + '0', '(-g1(x, y)/h_x + g1(x + h_x, y)/h_x)*f1(x, y) - g1(x, y)' + )), + ('Eq(f1 * g1.dx, g1).evaluate', 'f1', + ( + 'g1(x, y)', '(-g1(x, y)/h_x + g1(x + h_x, y)/h_x)*f1(x, y)' + )), + ('Eq((f1 * g1.dx).dy, f1).evaluate', 'f1', + ( + '0', '(-1/h_y)*(-g1(x, y)/h_x + g1(x + h_x, y)/h_x)*f1(x, y) ' + '+ (-g1(x, y + h_y)/h_x + g1(x + h_x, y + h_y)/h_x)*f1(x, y + h_y)/h_y ' + '- f1(x, y)' + )), + ('Eq((f1 * g1.dx).dy, f1).evaluate', 'g1', + ( + 'f1(x, y)', '(-1/h_y)*(-g1(x, y)/h_x + g1(x + h_x, y)/h_x)*f1(x, y) + ' + '(-g1(x, y + h_y)/h_x + g1(x + h_x, y + h_y)/h_x)*f1(x, y + h_y)/h_y' + )), + ('Eq(f2.laplace, g2).evaluate', 'g2', + ( + '-(-2.0*f2(t, x, y)/h_x**2 + f2(t, x - h_x, y)/h_x**2 + f2(t, x + h_x, y)' + '/h_x**2) - (-2.0*f2(t, x, y)/h_y**2 + f2(t, x, y - h_y)/h_y**2 + ' + 'f2(t, x, y + h_y)/h_y**2)', '-g2(t, x, y)' + )), + ('Eq(f2.laplace, g2).evaluate', 'f2', + ( + 'g2(t, x, y)', '-2.0*f2(t, x, y)/h_x**2 + f2(t, x - h_x, y)/h_x**2 + ' + 'f2(t, x + h_x, y)/h_x**2 - 2.0*f2(t, x, y)/h_y**2 + f2(t, x, y - h_y)' + '/h_y**2 + f2(t, x, y + h_y)/h_y**2' + )), + ('Eq(f2.laplace, f2).evaluate', 'f2', + ( + '0', '-2.0*f2(t, x, y)/h_x**2 + f2(t, x - h_x, y)/h_x**2 + ' + 'f2(t, x + h_x, y)/h_x**2 - 2.0*f2(t, x, y)/h_y**2 + f2(t, x, y - h_y)/h_y**2' + ' + f2(t, x, y + h_y)/h_y**2 - f2(t, x, y)' + )), + ('Eq(g2*f2.laplace, f2).evaluate', 'g2', + ( + 'f2(t, x, y)', '(-2.0*f2(t, x, y)/h_x**2 + f2(t, x - h_x, y)/h_x**2 + ' + 'f2(t, x + h_x, y)/h_x**2 - 2.0*f2(t, x, y)/h_y**2 + f2(t, x, y - h_y)/h_y**2' + ' + f2(t, x, y + h_y)/h_y**2)*g2(t, x, y)' + )), + ('Eq(f2.forward.laplace, f2).evaluate', 'f2.forward', + ( + 'f2(t, x, y)', '-2.0*f2(t + dt, x, y)/h_x**2 + f2(t + dt, x - h_x, y)/h_x**2' + ' + f2(t + dt, x + h_x, y)/h_x**2 - 2.0*f2(t + dt, x, y)/h_y**2 + ' + 'f2(t + dt, x, y - h_y)/h_y**2 + f2(t + dt, x, y + h_y)/h_y**2' + )), + ('Eq(f2.forward.laplace, f2).evaluate', 'f2', + ( + '-(-2.0*f2(t + dt, x, y)/h_x**2 + f2(t + dt, x - h_x, y)/h_x**2 + ' + 'f2(t + dt, x + h_x, y)/h_x**2) - (-2.0*f2(t + dt, x, y)/h_y**2 + ' + 'f2(t + dt, x, y - h_y)/h_y**2 + f2(t + dt, x, y + h_y)/h_y**2)', + '-f2(t, x, y)' + )), + ('Eq(f2.laplace + f2.forward.laplace, g2).evaluate', 'f2.forward', + ( + '-(-2.0*f2(t, x, y)/h_x**2 + f2(t, x - h_x, y)/h_x**2 + f2(t, x + h_x, y)/' + 'h_x**2) - (-2.0*f2(t, x, y)/h_y**2 + f2(t, x, y - h_y)/h_y**2 + ' + 'f2(t, x, y + h_y)/h_y**2) + g2(t, x, y)', '-2.0*f2(t + dt, x, y)/h_x**2 + ' + 'f2(t + dt, x - h_x, y)/h_x**2 + f2(t + dt, x + h_x, y)/h_x**2 - 2.0*' + 'f2(t + dt, x, y)/h_y**2 + f2(t + dt, x, y - h_y)/h_y**2 + ' + 'f2(t + dt, x, y + h_y)/h_y**2' + )), + ('Eq(g2.laplace, f2 + g2.forward).evaluate', 'g2.forward', + ( + '-(-2.0*g2(t, x, y)/h_x**2 + g2(t, x - h_x, y)/h_x**2 + ' + 'g2(t, x + h_x, y)/h_x**2) - (-2.0*g2(t, x, y)/h_y**2 + g2(t, x, y - h_y)' + '/h_y**2 + g2(t, x, y + h_y)/h_y**2) + f2(t, x, y)', '-g2(t + dt, x, y)' + )) +]) +def test_separate_eval_eqn(eqn, target, expected): + """ + Test the separate_eqn function on pre-evaluated equations. + """ + grid = Grid((2, 2)) + + so = 2 + + f1 = Function(name='f1', grid=grid, space_order=so) # noqa + g1 = Function(name='g1', grid=grid, space_order=so) # noqa + + f2 = TimeFunction(name='f2', grid=grid, space_order=so) # noqa + g2 = TimeFunction(name='g2', grid=grid, space_order=so) # noqa + + b, F, _, _ = separate_eqn(eval(eqn), eval(target)) + expected_b, expected_F = expected + + assert str(b) == expected_b + assert str(F) == expected_F + + +@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 + + f2 = TimeFunction(name='f2', grid=grid, space_order=so) # noqa + g2 = TimeFunction(name='g2', grid=grid, space_order=so) # noqa + + centre = centre_stencil(eval(expr), eval(target)) + + assert str(centre) == expected From 3000bd8e7d7541c684897324e9f0c4d2287d46e0 Mon Sep 17 00:00:00 2001 From: ZoeLeibowitz Date: Tue, 3 Jun 2025 22:47:00 +0100 Subject: [PATCH 35/84] types: Move EssentialBC to another file --- devito/petsc/__init__.py | 1 + devito/petsc/solve.py | 80 +++++++++++----------------------- devito/petsc/types/__init__.py | 2 + 3 files changed, 29 insertions(+), 54 deletions(-) diff --git a/devito/petsc/__init__.py b/devito/petsc/__init__.py index 2927bc960e..6aab617e8d 100644 --- a/devito/petsc/__init__.py +++ b/devito/petsc/__init__.py @@ -1 +1,2 @@ from devito.petsc.solve import * # noqa +from devito.petsc.types.equation import * # noqa \ No newline at end of file diff --git a/devito/petsc/solve.py b/devito/petsc/solve.py index f3e7cd377b..1004a403b6 100644 --- a/devito/petsc/solve.py +++ b/devito/petsc/solve.py @@ -16,9 +16,10 @@ from devito.tools import as_tuple, filter_ordered from devito.petsc.types import (LinearSolveExpr, PETScArray, DMDALocalInfo, FieldData, MultipleFieldData, Jacobian) +from devito.petsc.types.equation import EssentialBC, ZeroRow, ZeroColumn -__all__ = ['PETScSolve', 'EssentialBC'] +__all__ = ['PETScSolve'] def PETScSolve(target_eqns, target=None, solver_parameters=None, **kwargs): @@ -329,67 +330,38 @@ def generate_arrays_combined(self, *targets): } -class EssentialBC(Eq): - """ - A special equation used to handle essential boundary conditions - in the PETSc solver. Until PetscSection + DMDA is supported, - we treat essential boundary conditions as trivial equations - in the solver, where we place 1.0 (scaled) on the diagonal of - the jacobian, zero symmetrically and move the boundary - data to the right-hand side. - - NOTE: When users define essential boundary conditions, they need to ensure that - the SubDomains do not overlap. Solver will still run but may see unexpected behaviour - at boundaries. This will be documented in the PETSc examples. - """ - pass - - -class ZeroRow(EssentialBC): - """ - Equation used to zero the row of the Jacobian corresponding - to an essential BC. - This is only used interally by the compiler, not by users. - """ - pass - - -class ZeroColumn(EssentialBC): - """ - Equation used to zero the column of the Jacobian corresponding - to an essential BC. - This is only used interally by the compiler, not by users. - """ - pass +# class EssentialBC(Eq): +# """ +# A special equation used to handle essential boundary conditions +# in the PETSc solver. Until PetscSection + DMDA is supported, +# we treat essential boundary conditions as trivial equations +# in the solver, where we place 1.0 (scaled) on the diagonal of +# the jacobian, zero symmetrically and move the boundary +# data to the right-hand side. + +# NOTE: When users define essential boundary conditions, they need to ensure that +# the SubDomains do not overlap. Solver will still run but may see unexpected behaviour +# at boundaries. This will be documented in the PETSc examples. +# """ +# pass -# def separate_eqn(eqn, target): +# class ZeroRow(EssentialBC): # """ -# Separate the equation into two separate expressions, -# where F(target) = b. +# Equation used to zero the row of the Jacobian corresponding +# to an essential BC. +# This is only used interally by the compiler, not by users. # """ -# zeroed_eqn = Eq(eqn.lhs - eqn.rhs, 0) -# zeroed_eqn = eval_time_derivatives(zeroed_eqn.lhs) -# target_funcs = set(generate_targets(zeroed_eqn, target)) -# b, F_target = remove_targets(zeroed_eqn, target_funcs) -# return -b, F_target, zeroed_eqn, target_funcs +# pass -# def generate_targets(eq, target): +# class ZeroColumn(EssentialBC): # """ -# Extract all the functions that share the same time index as the target -# but may have different spatial indices. +# Equation used to zero the column of the Jacobian corresponding +# to an essential BC. +# This is only used interally by the compiler, not by users. # """ -# funcs = retrieve_functions(eq) -# if isinstance(target, TimeFunction): -# time_idx = target.indices[target.time_dim] -# targets = [ -# f for f in funcs if f.function is target.function and time_idx -# in f.indices -# ] -# else: -# targets = [f for f in funcs if f.function is target.function] -# return targets +# pass def targets_to_arrays(array, targets): diff --git a/devito/petsc/types/__init__.py b/devito/petsc/types/__init__.py index ebcceb8d45..f2305a8352 100644 --- a/devito/petsc/types/__init__.py +++ b/devito/petsc/types/__init__.py @@ -1,3 +1,5 @@ from .array import * # noqa from .types import * # noqa from .object import * # noqa +from .equation import * # noqa +from .macros import * # noqa From a0457f057645cd8b034c01d5accb9b233316b2ea Mon Sep 17 00:00:00 2001 From: ZoeLeibowitz Date: Tue, 3 Jun 2025 22:48:24 +0100 Subject: [PATCH 36/84] dsl/compiler: Add equation.py to petsc module --- devito/petsc/types/equation.py | 38 ++++++++++++++++++++++++++++++++++ 1 file changed, 38 insertions(+) create mode 100644 devito/petsc/types/equation.py diff --git a/devito/petsc/types/equation.py b/devito/petsc/types/equation.py new file mode 100644 index 0000000000..a5c5bc55ef --- /dev/null +++ b/devito/petsc/types/equation.py @@ -0,0 +1,38 @@ +from devito.types.equation import Eq + + +__all__ = ['EssentialBC'] + + +class EssentialBC(Eq): + """ + A special equation used to handle essential boundary conditions + in the PETSc solver. Until PetscSection + DMDA is supported, + we treat essential boundary conditions as trivial equations + in the solver, where we place 1.0 (scaled) on the diagonal of + the jacobian, zero symmetrically and move the boundary + data to the right-hand side. + + NOTE: When users define essential boundary conditions, they need to ensure that + the SubDomains do not overlap. Solver will still run but may see unexpected behaviour + at boundaries. This will be documented in the PETSc examples. + """ + pass + + +class ZeroRow(EssentialBC): + """ + Equation used to zero the row of the Jacobian corresponding + to an essential BC. + This is only used interally by the compiler, not by users. + """ + pass + + +class ZeroColumn(EssentialBC): + """ + Equation used to zero the column of the Jacobian corresponding + to an essential BC. + This is only used interally by the compiler, not by users. + """ + pass \ No newline at end of file From 7132eb113dd479ae95d2b67324a95606511735fe Mon Sep 17 00:00:00 2001 From: ZoeLeibowitz Date: Wed, 4 Jun 2025 10:55:51 +0100 Subject: [PATCH 37/84] dsl: Add jacobian class for single fields --- devito/petsc/iet/routines.py | 9 +- devito/petsc/solve.py | 260 ++++---------------------------- devito/petsc/types/types.py | 278 ++++++++++++++++++++++++++++++----- devito/types/grid.py | 7 +- 4 files changed, 280 insertions(+), 274 deletions(-) diff --git a/devito/petsc/iet/routines.py b/devito/petsc/iet/routines.py index 00c46f8b7d..1fc4cd7758 100644 --- a/devito/petsc/iet/routines.py +++ b/devito/petsc/iet/routines.py @@ -108,7 +108,8 @@ def zero_memory(self): def _make_core(self): fielddata = self.injectsolve.expr.rhs.fielddata - self._make_matvec(fielddata.arrays, fielddata.matvecs) + # from IPython.core.debugger import set_trace + self._make_matvec(fielddata.arrays, fielddata.jacobian.matvecs) self._make_formfunc(fielddata) self._make_formrhs(fielddata) if fielddata.initialguess: @@ -266,7 +267,7 @@ def _create_matvec_body(self, body, arrays): return matvec_body def _make_formfunc(self, fielddata): - formfuncs = fielddata.formfuncs + formfuncs = fielddata.residual.formfuncs # Compile formfunc `eqns` into an IET via recursive compilation irs_formfunc, _ = self.rcompile( formfuncs, options={'mpi': False}, sregistry=self.sregistry, @@ -403,7 +404,7 @@ def _create_formfunc_body(self, body, fielddata): return Uxreplace(subs).visit(formfunc_body) def _make_formrhs(self, fielddata): - formrhs = fielddata.formrhs + formrhs = fielddata.residual.formrhs sobjs = self.solver_objs # Compile formrhs `eqns` into an IET via recursive compilation @@ -758,7 +759,7 @@ def _whole_matvec_body(self): ) def _make_whole_formfunc(self, fielddata): - formfuncs = fielddata.formfuncs + formfuncs = fielddata.residual.formfuncs # Compile formfunc `eqns` into an IET via recursive compilation irs_formfunc, _ = self.rcompile( formfuncs, options={'mpi': False}, sregistry=self.sregistry, diff --git a/devito/petsc/solve.py b/devito/petsc/solve.py index 1004a403b6..eb8f171b21 100644 --- a/devito/petsc/solve.py +++ b/devito/petsc/solve.py @@ -15,7 +15,8 @@ generate_targets) from devito.tools import as_tuple, filter_ordered from devito.petsc.types import (LinearSolveExpr, PETScArray, DMDALocalInfo, - FieldData, MultipleFieldData, Jacobian) + FieldData, MultipleFieldData, Jacobian, Residual, + MixedResidual, MixedJacobian) from devito.petsc.types.equation import EssentialBC, ZeroRow, ZeroColumn @@ -26,7 +27,7 @@ def PETScSolve(target_eqns, target=None, solver_parameters=None, **kwargs): if target is not None: return InjectSolve(solver_parameters, {target: target_eqns}).build_eq() else: - return InjectSolveNested(solver_parameters, target_eqns).build_eq() + return InjectMixedSolve(solver_parameters, target_eqns).build_eq() class InjectSolve: @@ -34,14 +35,6 @@ def __init__(self, solver_parameters=None, target_eqns=None): self.solver_params = solver_parameters self.time_mapper = None self.target_eqns = target_eqns - # TODO: make this _ - self.cell_area = None - # self._centre_stencils = set() - self._diag_scale = defaultdict(set) - - @property - def diag_scale(self): - return self._diag_scale def build_eq(self): target, funcs, fielddata = self.linear_solve_args() @@ -59,42 +52,20 @@ def build_eq(self): def linear_solve_args(self): target, eqns = next(iter(self.target_eqns.items())) eqns = as_tuple(eqns) - self.cell_area = np.prod(target.grid.spacing_symbols) funcs = get_funcs(eqns) self.time_mapper = generate_time_mapper(funcs) arrays = self.generate_arrays(target) - return target, tuple(funcs), self.generate_field_data(eqns, target, arrays) - - def generate_field_data(self, eqns, target, arrays): - # Apply essential boundary conditions first to preserve - # operator symmetry during Jacobian "construction" eqns = sorted(eqns, key=lambda e: 0 if isinstance(e, EssentialBC) else 1) - matvecs = [e for eq in eqns for e in self.build_matvec_eq(eq, target, arrays)] - - formfuncs, formrhs = map( - lambda x: [e for i in x for e in i], - zip(*[self.build_function_eq(eq, target, arrays) for eq in eqns]) - ) - - stencils = set() - for eq in matvecs: - if not isinstance(eq, EssentialBC): - stencil = centre_stencil(eq.rhs, arrays['x'], as_coeff=True) - stencils.add(stencil) + jacobian = Jacobian(target, self.time_mapper, arrays) + jacobian.build_block(eqns) - if len(stencils) > 1: - # Scaling of jacobian is therefore ambiguous, potentially could average across the subblock - # for now just set to trivial 1.0 - scale = 1.0 - else: - scale = next(iter(stencils)) + scale = 1.0 - # from IPython import embed; embed() - matvecs = self.scale_essential_bcs(matvecs, scale) - formfuncs = self.scale_essential_bcs(formfuncs, scale) + residual = Residual(target, self.time_mapper, arrays, scale) + residual.build_equations(eqns) initialguess = [ eq for eq in @@ -102,70 +73,15 @@ def generate_field_data(self, eqns, target, arrays): if eq is not None ] - return FieldData( + field_data = FieldData( target=target, - matvecs=matvecs, - formfuncs=formfuncs, - formrhs=formrhs, + jacobian=jacobian, + residual=residual, initialguess=initialguess, arrays=arrays ) - def build_function_eq(self, eq, target, arrays): - b, F_target, _, targets = separate_eqn(eq, target) - formfunc = self.make_formfunc(eq, F_target, arrays, targets) - formrhs = self.make_rhs(eq, b, arrays) - - return (formfunc, formrhs) - - def build_matvec_eq(self, eq, target, arrays): - b, F_target, _, targets = separate_eqn(eq, target) - if F_target: - return self.make_matvec(eq, F_target, arrays, targets) - return (None,) - - def make_matvec(self, eq, F_target, arrays, targets): - if isinstance(eq, EssentialBC): - # NOTE: Until PetscSection + DMDA is supported, we leave - # the essential BCs in the solver. - # Trivial equations for bc rows -> place 1.0 on diagonal (scaled) - # and zero symmetrically. - rhs = arrays['x'] - zero_row = ZeroRow(arrays['y'], rhs, subdomain=eq.subdomain) - zero_column = ZeroColumn(arrays['x'], 0.0, subdomain=eq.subdomain) - return (zero_row, zero_column) - else: - rhs = F_target.subs(targets_to_arrays(arrays['x'], targets)) - rhs = rhs.subs(self.time_mapper) * self.cell_area - # TODO: Average centre stencils if they vary, to scale essential BC rows. - # self.centre = centre_stencil(rhs, arrays['x'], as_coeff=True) - # stencil = centre_stencil(rhs, arrays['x'], as_coeff=True) - # self._centre_stencils[arrays['x']].add(stencil) - # self._centre_stencils.add(stencil) - - return as_tuple(Eq(arrays['y'], rhs, subdomain=eq.subdomain)) - - def make_formfunc(self, eq, F_target, arrays, targets): - if isinstance(eq, EssentialBC): - # The initial guess satisfies the essential BCs, so this term is zero. - # Still included to support Jacobian testing via finite differences. - rhs = arrays['x'] - eq.rhs - zero_row = ZeroRow(arrays['f'], rhs, subdomain=eq.subdomain) - # Move essential boundary condition to the right-hand side - zero_col = ZeroColumn(arrays['x'], eq.rhs, subdomain=eq.subdomain) - return (zero_row, zero_col) - else: - if isinstance(F_target, (int, float)): - rhs = F_target * self.cell_area - else: - rhs = F_target.subs(targets_to_arrays(arrays['x'], targets)) - rhs = rhs.subs(self.time_mapper) * self.cell_area - return as_tuple(Eq(arrays['f'], rhs, subdomain=eq.subdomain)) - - def make_rhs(self, eq, b, arrays): - rhs = 0. if isinstance(eq, EssentialBC) else b.subs(self.time_mapper) - rhs = rhs * self.cell_area - return as_tuple(Eq(arrays['b'], rhs, subdomain=eq.subdomain)) + return target, tuple(funcs), field_data def make_initial_guess(self, eq, target, arrays): """ @@ -192,25 +108,8 @@ def generate_arrays(self, target): for p in prefixes } - def scale_essential_bcs(self, equations, scale): - """ - Scale the essential boundary rows so that the Jacobian has a constant diagonal, - thereby reducing its condition number. - """ - # # stencils = self.centre_stencils[arrays['x']] - # if len(stencils) > 1: - # # Scaling of jacobian is therefore ambiguous, potentially could averge across the subblock - # # for now just set to trivial 1.0 - # scale = 1.0 - # else: - # scale = next(iter(stencils)) - return [ - eq._rebuild(rhs=scale * eq.rhs) if isinstance(eq, ZeroRow) else eq - for eq in equations - ] - -class InjectSolveNested(InjectSolve): +class InjectMixedSolve(InjectSolve): def linear_solve_args(self): @@ -221,99 +120,24 @@ def linear_solve_args(self): self.time_mapper = generate_time_mapper(funcs) coupled_targets = list(self.target_eqns.keys()) - jacobian = Jacobian(coupled_targets) arrays = self.generate_arrays_combined(*coupled_targets) - all_data = MultipleFieldData(jacobian=jacobian, arrays=arrays, - targets=coupled_targets) - - self.cell_area = np.prod(all_data.grid.spacing_symbols) + jacobian = MixedJacobian(coupled_targets, self.time_mapper, arrays) + jacobian.build_blocks(self.target_eqns) - all_formfuncs = [] - - for target, eqns in self.target_eqns.items(): - - # Update all rows of the Jacobian for this target - self.update_jacobian(as_tuple(eqns), target, jacobian, arrays[target]) - - formfuncs = chain.from_iterable( - self.build_function_eq(eq, target, coupled_targets, arrays) - for eq in as_tuple(eqns) - ) - # from IPython import embed; embed() - scale, = self._diag_scale[arrays[target]['x']] - all_formfuncs.extend(self.scale_essential_bcs(formfuncs, scale)) - - formfuncs = tuple(sorted( - all_formfuncs, key=lambda e: not isinstance(e, EssentialBC) - )) - all_data.extend_formfuncs(formfuncs) - - return target, tuple(funcs), all_data - - def update_jacobian(self, eqns, target, jacobian, arrays): - - for submat, mtvs in jacobian.submatrices[target].items(): - matvecs = [ - e for eq in eqns for e in - self.build_matvec_eq(eq, mtvs['derivative_wrt'], arrays) - ] - matvecs = [m for m in matvecs if m is not None] - - if submat in jacobian.diagonal_submatrix_keys: - stencils = set() - for eq in matvecs: - if not isinstance(eq, EssentialBC): - stencil = centre_stencil(eq.rhs, arrays['x'], as_coeff=True) - stencils.add(stencil) - # from IPython import embed; embed() - if len(stencils) > 1: - # Scaling of jacobian is therefore ambiguous, potentially could average across the subblock - # for now just set to trivial 1.0 - # TODO: doens't need to be a defaultdict, just a dict? - self._diag_scale[arrays['x']].add(1.0) - scale = 1.0 - else: - scale = next(iter(stencils)) - self._diag_scale[arrays['x']].add(scale) - # scale = next(iter(stencils)) - - matvecs = self.scale_essential_bcs(matvecs, scale) - - matvecs = tuple(sorted(matvecs, key=lambda e: not isinstance(e, EssentialBC))) - - if matvecs: - jacobian.set_submatrix(target, submat, matvecs) - - def build_function_eq(self, eq, main_target, coupled_targets, arrays): - zeroed = eq.lhs - eq.rhs - - zeroed_eqn = Eq(eq.lhs - eq.rhs, 0) - eval_zeroed_eqn = eval_time_derivatives(zeroed_eqn.lhs) - - mapper = {} - for t in coupled_targets: - target_funcs = set(generate_targets(Eq(eval_zeroed_eqn, 0), t)) - mapper.update(targets_to_arrays(arrays[t]['x'], target_funcs)) + residual = MixedResidual(coupled_targets, self.time_mapper, arrays) + residual.build_equations(self.target_eqns) - if isinstance(eq, EssentialBC): - rhs = arrays[main_target]['x'] - eq.rhs - zero_row = ZeroRow( - arrays[main_target]['f'], rhs, subdomain=eq.subdomain - ) - zero_col = ZeroColumn( - arrays[main_target]['x'], eq.rhs, subdomain=eq.subdomain - ) - return (zero_row, zero_col) - else: - if isinstance(zeroed, (int, float)): - rhs = zeroed * self.cell_area - else: - rhs = zeroed.subs(mapper) - rhs = rhs.subs(self.time_mapper)*self.cell_area + all_data = MultipleFieldData( + targets=coupled_targets, + arrays=arrays, + jacobian=jacobian, + residual=residual + ) - return as_tuple(Eq(arrays[main_target]['f'], rhs, subdomain=eq.subdomain)) + # TODO: rethink what target to return here??? + return coupled_targets[0], tuple(funcs), all_data def generate_arrays_combined(self, *targets): return { @@ -330,40 +154,6 @@ def generate_arrays_combined(self, *targets): } -# class EssentialBC(Eq): -# """ -# A special equation used to handle essential boundary conditions -# in the PETSc solver. Until PetscSection + DMDA is supported, -# we treat essential boundary conditions as trivial equations -# in the solver, where we place 1.0 (scaled) on the diagonal of -# the jacobian, zero symmetrically and move the boundary -# data to the right-hand side. - -# NOTE: When users define essential boundary conditions, they need to ensure that -# the SubDomains do not overlap. Solver will still run but may see unexpected behaviour -# at boundaries. This will be documented in the PETSc examples. -# """ -# pass - - -# class ZeroRow(EssentialBC): -# """ -# Equation used to zero the row of the Jacobian corresponding -# to an essential BC. -# This is only used interally by the compiler, not by users. -# """ -# pass - - -# class ZeroColumn(EssentialBC): -# """ -# Equation used to zero the column of the Jacobian corresponding -# to an essential BC. -# This is only used interally by the compiler, not by users. -# """ -# pass - - def targets_to_arrays(array, targets): """ Map each target in `targets` to a corresponding array generated from `array`, diff --git a/devito/petsc/types/types.py b/devito/petsc/types/types.py index 75ee8e8254..3897552f7b 100644 --- a/devito/petsc/types/types.py +++ b/devito/petsc/types/types.py @@ -1,8 +1,14 @@ import sympy +from itertools import chain + from devito.tools import Reconstructable, sympy_mutex, as_tuple from devito.tools.dtypes_lowering import dtype_mapper from devito.petsc.utils import petsc_variables +from devito.symbolics.extraction import separate_eqn, generate_targets +from devito.petsc.types.equation import EssentialBC, ZeroRow, ZeroColumn +from devito.types.equation import Eq +from devito.operations.solve import eval_time_derivatives class MetaData(sympy.Function, Reconstructable): @@ -131,7 +137,7 @@ def eval(cls, *args): class FieldData: - def __init__(self, target=None, matvecs=None, formfuncs=None, formrhs=None, + def __init__(self, target=None, jacobian=None, residual=None, initialguess=None, arrays=None, **kwargs): self._target = target petsc_precision = dtype_mapper[petsc_variables['PETSC_PRECISION']] @@ -141,9 +147,8 @@ def __init__(self, target=None, matvecs=None, formfuncs=None, formrhs=None, f"PETSc configuration. " f"Expected {petsc_precision}, but got {self._target.dtype}." ) - self._matvecs = matvecs - self._formfuncs = formfuncs - self._formrhs = formrhs + self._jacobian = jacobian + self._residual = residual self._initialguess = initialguess self._arrays = arrays @@ -152,16 +157,12 @@ def target(self): return self._target @property - def matvecs(self): - return self._matvecs - - @property - def formfuncs(self): - return self._formfuncs + def jacobian(self): + return self._jacobian @property - def formrhs(self): - return self._formrhs + def residual(self): + return self._residual @property def initialguess(self): @@ -189,14 +190,11 @@ def targets(self): class MultipleFieldData(FieldData): - def __init__(self, targets, arrays, jacobian=None): + def __init__(self, targets, arrays, jacobian=None, residual=None): self._targets = as_tuple(targets) self._arrays = arrays self._jacobian = jacobian - self._formfuncs = [] - - def extend_formfuncs(self, formfuncs): - self._formfuncs.extend(formfuncs) + self._residual = residual @property def space_dimensions(self): @@ -239,14 +237,13 @@ def jacobian(self): def targets(self): return self._targets - @property - def arrays(self): - return self._arrays -class Jacobian: - def __init__(self, targets): - self.targets = targets +class BaseJacobian: + def __init__(self, targets, time_mapper, arrays): + self.targets = as_tuple(targets) + self.time_mapper = time_mapper + self.arrays = arrays self.submatrices = self._initialize_submatrices() def _initialize_submatrices(self): @@ -298,6 +295,18 @@ def submat_to_index(self): for key, value in submats.items() } + @property + def diagonal_submatrix_keys(self): + """ + Return a list of diagonal submatrix keys (e.g., ['J00', 'J11']). + """ + keys = [] + for i, target in enumerate(self.targets): + diag_key = f'J{i}{i}' + if diag_key in self.submatrices[target]: + keys.append(diag_key) + return keys + def set_submatrix(self, field, key, matvecs): """ Set a specific submatrix for a field. @@ -321,18 +330,219 @@ def get_submatrix(self, field, key): Retrieve a specific submatrix. """ return self.submatrices.get(field, {}).get(key, None) - + + def build_matvec_eq(self, eq, target, arrays): + b, F_target, _, targets = separate_eqn(eq, target) + if F_target: + return self.make_matvec(eq, F_target, targets, arrays) + return (None,) + + def make_matvec(self, eq, F_target, targets, arrays): + if isinstance(eq, EssentialBC): + # NOTE: Until PetscSection + DMDA is supported, we leave + # the essential BCs in the solver. + # Trivial equations for bc rows -> place 1.0 on diagonal (scaled) + # and zero symmetrically. + rhs = arrays['x'] + zero_row = ZeroRow(arrays['y'], rhs, subdomain=eq.subdomain) + zero_column = ZeroColumn(arrays['x'], 0.0, subdomain=eq.subdomain) + return (zero_row, zero_column) + else: + rhs = F_target.subs(targets_to_arrays(arrays['x'], targets)) + # rhs = rhs.subs(self.time_mapper) * self.cell_area + rhs = rhs = rhs.subs(self.time_mapper) + + return as_tuple(Eq(arrays['y'], rhs, subdomain=eq.subdomain)) + + def __repr__(self): + return str(self.submatrices) + + +class Jacobian(BaseJacobian): + @property - def diagonal_submatrix_keys(self): + def target(self): + return self.targets[0] + + @property + def matvecs(self): + return self.submatrices[self.target]['J00']['matvecs'] + + # TODO: use same structure arrays for both jacobian and mixedjacobian + def build_block(self, eqns): + for submat, mtvs in self.submatrices[self.target].items(): + matvecs = [ + e for eq in eqns for e in + self.build_matvec_eq(eq, mtvs['derivative_wrt'], self.arrays) + ] + matvecs = [m for m in matvecs if m is not None] + matvecs = tuple(sorted(matvecs, key=lambda e: not isinstance(e, EssentialBC))) + + if matvecs: + self.set_submatrix(self.target, submat, matvecs) + + +class MixedJacobian(BaseJacobian): + + # TODO: use same structure arrays for both jacobian and mixedjacobian + def build_block(self, target, eqns): + for submat, mtvs in self.submatrices[target].items(): + matvecs = [ + e for eq in eqns for e in + self.build_matvec_eq(eq, mtvs['derivative_wrt'], self.arrays[target]) + ] + matvecs = [m for m in matvecs if m is not None] + matvecs = tuple(sorted(matvecs, key=lambda e: not isinstance(e, EssentialBC))) + + if matvecs: + self.set_submatrix(target, submat, matvecs) + + def build_blocks(self, target_eqns): + for target, eqns in target_eqns.items(): + self.build_block(target, eqns) + + +class BaseResidual: + def scale_essential_bcs(self, equations): """ - Return a list of diagonal submatrix keys (e.g., ['J00', 'J11']). """ - keys = [] - for i, target in enumerate(self.targets): - diag_key = f'J{i}{i}' - if diag_key in self.submatrices[target]: - keys.append(diag_key) - return keys + return [ + eq._rebuild(rhs=self.scale * eq.rhs) if isinstance(eq, ZeroRow) else eq + for eq in equations + ] - def __repr__(self): - return str(self.submatrices) + +class Residual(BaseResidual): + """ + """ + + def __init__(self, target, time_mapper, arrays, scale): + self.target = target + self.time_mapper = time_mapper + self.arrays = arrays + self.scale = scale + self.formfuncs = [] + self.formrhs = [] + + def build_equations(self, eqns): + """ + """ + for eq in eqns: + b, F_target, _, targets = separate_eqn(eq, self.target) + F_target = self.make_F_target(eq, F_target, targets) + b = self.make_b(eq, b) + self.formfuncs.extend(F_target) + self.formrhs.extend(b) + + self.formfuncs = self.scale_essential_bcs(self.formfuncs) + + def make_F_target(self, eq, F_target, targets): + arrays = self.arrays + volume = self.target.grid.symbolic_volume_cell + if isinstance(eq, EssentialBC): + # The initial guess satisfies the essential BCs, so this term is zero. + # Still included to support Jacobian testing via finite differences. + rhs = arrays['x'] - eq.rhs + zero_row = ZeroRow(arrays['f'], rhs, subdomain=eq.subdomain) + # Move essential boundary condition to the right-hand side + zero_col = ZeroColumn(arrays['x'], eq.rhs, subdomain=eq.subdomain) + return (zero_row, zero_col) + else: + if isinstance(F_target, (int, float)): + rhs = F_target * volume + else: + rhs = F_target.subs(targets_to_arrays(arrays['x'], targets)) + rhs = rhs.subs(self.time_mapper) * volume + return as_tuple(Eq(arrays['f'], rhs, subdomain=eq.subdomain)) + + def make_b(self, eq, b): + rhs = 0. if isinstance(eq, EssentialBC) else b.subs(self.time_mapper) + rhs = rhs * self.target.grid.symbolic_volume_cell + return as_tuple(Eq(self.arrays['b'], rhs, subdomain=eq.subdomain)) + + +class MixedResidual(BaseResidual): + """ + """ + # TODO: change default and pass in correct scale + def __init__(self, targets, time_mapper, arrays, scale=1.0): + self.targets = as_tuple(targets) + self.time_mapper = time_mapper + self.arrays = arrays + self.scale = scale + self.formfuncs = [] + + def build_equations(self, eqn_dict): + all_formfuncs = [] + for target, eqns in eqn_dict.items(): + + formfuncs = chain.from_iterable( + self.build_function_eq(eq, target) + for eq in as_tuple(eqns) + ) + + # scale, = self._diag_scale[arrays[target]['x']] + # fix this + # scale = 1.0 + all_formfuncs.extend(self.scale_essential_bcs(formfuncs)) + + self.formfuncs = tuple(sorted( + all_formfuncs, key=lambda e: not isinstance(e, EssentialBC) + )) + + + def build_function_eq(self, eq, target): + zeroed = eq.lhs - eq.rhs + + zeroed_eqn = Eq(eq.lhs - eq.rhs, 0) + eval_zeroed_eqn = eval_time_derivatives(zeroed_eqn.lhs) + + mapper = {} + for t in self.targets: + target_funcs = set(generate_targets(Eq(eval_zeroed_eqn, 0), t)) + mapper.update(targets_to_arrays(self.arrays[t]['x'], target_funcs)) + + if isinstance(eq, EssentialBC): + rhs = self.arrays[target]['x'] - eq.rhs + zero_row = ZeroRow( + self.arrays[target]['f'], rhs, subdomain=eq.subdomain + ) + zero_col = ZeroColumn( + self.arrays[target]['x'], eq.rhs, subdomain=eq.subdomain + ) + return (zero_row, zero_col) + else: + if isinstance(zeroed, (int, float)): + # rhs = zeroed * self.cell_area + rhs = zeroed + else: + rhs = zeroed.subs(mapper) + # rhs = rhs.subs(self.time_mapper)*self.cell_area + rhs = rhs.subs(self.time_mapper) + + return as_tuple(Eq(self.arrays[target]['f'], rhs, subdomain=eq.subdomain)) + + + +def targets_to_arrays(array, targets): + """ + Map each target in `targets` to a corresponding array generated from `array`, + matching the spatial indices of the target. + Example: + -------- + >>> array + vec_u(x, y) + >>> targets + {u(t + dt, x + h_x, y), u(t + dt, x - h_x, y), u(t + dt, x, y)} + >>> targets_to_arrays(array, targets) + {u(t + dt, x - h_x, y): vec_u(x - h_x, y), + u(t + dt, x + h_x, y): vec_u(x + h_x, y), + u(t + dt, x, y): vec_u(x, y)} + """ + space_indices = [ + tuple(f.indices[d] for d in f.space_dimensions) for f in targets + ] + array_targets = [ + array.subs(dict(zip(array.indices, i))) for i in space_indices + ] + return dict(zip(targets, array_targets)) diff --git a/devito/types/grid.py b/devito/types/grid.py index f0d4a440d5..404619f3c0 100644 --- a/devito/types/grid.py +++ b/devito/types/grid.py @@ -280,10 +280,15 @@ def interior(self): """The interior SubDomain of the Grid.""" return self.subdomains['interior'] + @property + def symbolic_volume_cell(self): + """Symbolic volume of a single cell e.g. h_x*h_y*h_z in 3D.""" + return prod(d.spacing for d in self.dimensions) + @property def volume_cell(self): """Volume of a single cell e.g h_x*h_y*h_z in 3D.""" - return prod(d.spacing for d in self.dimensions).subs(self.spacing_map) + return self.symbolic_volume_cell.subs(self.spacing_map) @property def spacing(self): From 6897ab8490638586879bc7e7c907bd6f9a0abe87 Mon Sep 17 00:00:00 2001 From: ZoeLeibowitz Date: Wed, 4 Jun 2025 15:08:36 +0100 Subject: [PATCH 38/84] dsl: Fix compatible scaling for single fields --- .github/workflows/pytest-petsc.yml | 2 +- devito/petsc/iet/routines.py | 12 +- devito/petsc/solve.py | 136 ++--- devito/petsc/types/array.py | 4 + devito/petsc/types/equation.py | 39 +- devito/petsc/types/object.py | 1 - devito/petsc/types/types.py | 370 +++++++----- examples/petsc/Poisson/01_poisson.py | 7 +- examples/petsc/petsc_test.py | 2 +- .../random/{biharmonic => }/02_biharmonic.py | 0 .../random/biharmonic/biharmonic_matfree.c | 556 ------------------ .../biharmonic/biharmonic_matfree_nonscaled.c | 553 ----------------- tests/test_petsc.py | 1 - 13 files changed, 321 insertions(+), 1362 deletions(-) rename examples/petsc/random/{biharmonic => }/02_biharmonic.py (100%) delete mode 100644 examples/petsc/random/biharmonic/biharmonic_matfree.c delete mode 100644 examples/petsc/random/biharmonic/biharmonic_matfree_nonscaled.c diff --git a/.github/workflows/pytest-petsc.yml b/.github/workflows/pytest-petsc.yml index 19446642af..9c035a86e7 100644 --- a/.github/workflows/pytest-petsc.yml +++ b/.github/workflows/pytest-petsc.yml @@ -83,7 +83,7 @@ jobs: ${{ env.RUN_CMD }} mpiexec -n 1 python3 examples/petsc/Poisson/03_poisson.py ${{ env.RUN_CMD }} mpiexec -n 1 python3 examples/petsc/Poisson/04_poisson.py ${{ env.RUN_CMD }} mpiexec -n 1 python3 examples/petsc/random/01_helmholtz.py - ${{ env.RUN_CMD }} mpiexec -n 1 python3 examples/petsc/random/biharmonic/02_biharmonic.py + ${{ env.RUN_CMD }} mpiexec -n 1 python3 examples/petsc/random/02_biharmonic.py - name: Upload coverage to Codecov if: "!contains(matrix.name, 'docker')" diff --git a/devito/petsc/iet/routines.py b/devito/petsc/iet/routines.py index 1fc4cd7758..8f580a1b52 100644 --- a/devito/petsc/iet/routines.py +++ b/devito/petsc/iet/routines.py @@ -108,11 +108,10 @@ def zero_memory(self): def _make_core(self): fielddata = self.injectsolve.expr.rhs.fielddata - # from IPython.core.debugger import set_trace self._make_matvec(fielddata.arrays, fielddata.jacobian.matvecs) self._make_formfunc(fielddata) self._make_formrhs(fielddata) - if fielddata.initialguess: + if fielddata.initialguess.equations: self._make_initialguess(fielddata) self._make_user_struct_callback() @@ -516,7 +515,7 @@ def _create_form_rhs_body(self, body, fielddata): return Uxreplace(subs).visit(formrhs_body) def _make_initialguess(self, fielddata): - initguess = fielddata.initialguess + initguess = fielddata.initialguess.equations sobjs = self.solver_objs # Compile initital guess `eqns` into an IET via recursive compilation @@ -1208,7 +1207,6 @@ def _extend_build(self, base_dict): 'f': fbundle, 'x': xbundle, 'bundle_mapper': bundle_mapper, - # TODO: maybe this shouldn't be here 'target_indices': target_indices } @@ -1438,10 +1436,6 @@ def _create_dmda(self, dmda): class CoupledSetup(BaseSetup): - # @property - # def snes_ctx(self): - # return Byref(self.solver_objs['jacctx']) - def _setup(self): # TODO: minimise code duplication with superclass objs = self.objs @@ -1683,7 +1677,7 @@ def _execute_solve(self): struct_assignment = self.timedep.assign_time_iters(sobjs['userctx']) targets = self.injectsolve.expr.rhs.fielddata.targets - # TODO: Optimise the ccode generated here + # TODO: optimise the ccode generated here pre_solve = () post_solve = () diff --git a/devito/petsc/solve.py b/devito/petsc/solve.py index eb8f171b21..d2f0f39809 100644 --- a/devito/petsc/solve.py +++ b/devito/petsc/solve.py @@ -1,29 +1,55 @@ -from functools import singledispatch - -import sympy -import numpy as np -from itertools import chain -from collections import defaultdict - -from devito.finite_differences.differentiable import Mul -from devito.finite_differences.derivative import Derivative -from devito.types import Eq, Symbol, SteppingDimension, TimeFunction +from devito.types import Symbol, SteppingDimension from devito.types.equation import PetscEq from devito.operations.solve import eval_time_derivatives from devito.symbolics import retrieve_functions -from devito.symbolics.extraction import (separate_eqn, centre_stencil, - generate_targets) from devito.tools import as_tuple, filter_ordered from devito.petsc.types import (LinearSolveExpr, PETScArray, DMDALocalInfo, FieldData, MultipleFieldData, Jacobian, Residual, - MixedResidual, MixedJacobian) -from devito.petsc.types.equation import EssentialBC, ZeroRow, ZeroColumn + MixedResidual, MixedJacobian, InitialGuess) +from devito.petsc.types.equation import EssentialBC __all__ = ['PETScSolve'] -def PETScSolve(target_eqns, target=None, solver_parameters=None, **kwargs): +def PETScSolve(target_eqns, target=None, solver_parameters=None): + """ + Returns a symbolic equation representing a linear PETSc solver, + enriched with all the necessary metadata for execution within an `Operator`. + When passed to an `Operator`, this symbolic equation triggers code generation + and lowering to the PETSc backend. + + This function supports both single- and multi-target systems. In the multi-target + (mixed system) case, the solution vector spans all provided target fields. + + Parameters + ---------- + target_eqns : Eq or list of Eq, or dict of Function-like -> Eq or list of Eq + The targets and symbolic equations defining the system to be solved. + + - **Single-field problem**: + Pass a single Eq or list of Eq, and specify `target` separately: + PETScSolve(Eq1, target) + PETScSolve([Eq1, Eq2], target) + + - **Multi-field (mixed) problem**: + Pass a dictionary mapping each target field to its Eq(s): + PETScSolve({u: Eq1, v: Eq2}) + PETScSolve({u: [Eq1, Eq2], v: [Eq3, Eq4]}) + + target : Function-like + The function (e.g., `Function`, `TimeFunction`) into which the linear + system solves. This represents the solution vector updated by the solver. + + solver_parameters : dict, optional + PETSc solver options. + + Returns + ------- + Eq + A symbolic equation that wraps the system, solver metadata, + and boundary conditions. This can be passed directly to a Devito Operator. + """ if target is not None: return InjectSolve(solver_parameters, {target: target_eqns}).build_eq() else: @@ -59,19 +85,11 @@ def linear_solve_args(self): eqns = sorted(eqns, key=lambda e: 0 if isinstance(e, EssentialBC) else 1) - jacobian = Jacobian(target, self.time_mapper, arrays) - jacobian.build_block(eqns) - - scale = 1.0 + jacobian = Jacobian(target, eqns, arrays, self.time_mapper) - residual = Residual(target, self.time_mapper, arrays, scale) - residual.build_equations(eqns) + residual = Residual(target, eqns, arrays, self.time_mapper, jacobian.scdiag) - initialguess = [ - eq for eq in - (self.make_initial_guess(e, target, arrays) for e in eqns) - if eq is not None - ] + initialguess = InitialGuess(target, eqns, arrays) field_data = FieldData( target=target, @@ -83,22 +101,6 @@ def linear_solve_args(self): return target, tuple(funcs), field_data - def make_initial_guess(self, eq, target, arrays): - """ - Enforce initial guess to satisfy essential BCs. - # TODO: For time-stepping, only enforce these once outside the time loop - and use the previous time-step solution as the initial guess for next time step. - # TODO: Extend this to "coupled". - """ - if isinstance(eq, EssentialBC): - assert eq.lhs == target - return Eq( - arrays['x'], eq.rhs, - subdomain=eq.subdomain - ) - else: - return None - def generate_arrays(self, target): return { p: PETScArray(name=f'{p}_{target.name}', @@ -112,7 +114,6 @@ def generate_arrays(self, target): class InjectMixedSolve(InjectSolve): def linear_solve_args(self): - combined_eqns = [] for eqns in self.target_eqns.values(): combined_eqns.extend(eqns) @@ -123,11 +124,14 @@ def linear_solve_args(self): arrays = self.generate_arrays_combined(*coupled_targets) - jacobian = MixedJacobian(coupled_targets, self.time_mapper, arrays) - jacobian.build_blocks(self.target_eqns) + jacobian = MixedJacobian( + self.target_eqns, arrays, self.time_mapper + ) - residual = MixedResidual(coupled_targets, self.time_mapper, arrays) - residual.build_equations(self.target_eqns) + residual = MixedResidual( + self.target_eqns, arrays, + self.time_mapper, jacobian.target_scaler_mapper + ) all_data = MultipleFieldData( targets=coupled_targets, @@ -136,46 +140,10 @@ def linear_solve_args(self): residual=residual ) - # TODO: rethink what target to return here??? return coupled_targets[0], tuple(funcs), all_data def generate_arrays_combined(self, *targets): - return { - target: { - p: PETScArray( - name=f'{p}_{target.name}', - target=target, - liveness='eager', - localinfo=localinfo - ) - for p in prefixes - } - for target in targets - } - - -def targets_to_arrays(array, targets): - """ - Map each target in `targets` to a corresponding array generated from `array`, - matching the spatial indices of the target. - Example: - -------- - >>> array - vec_u(x, y) - >>> targets - {u(t + dt, x + h_x, y), u(t + dt, x - h_x, y), u(t + dt, x, y)} - >>> targets_to_arrays(array, targets) - {u(t + dt, x - h_x, y): vec_u(x - h_x, y), - u(t + dt, x + h_x, y): vec_u(x + h_x, y), - u(t + dt, x, y): vec_u(x, y)} - """ - space_indices = [ - tuple(f.indices[d] for d in f.space_dimensions) for f in targets - ] - array_targets = [ - array.subs(dict(zip(array.indices, i))) for i in space_indices - ] - return dict(zip(targets, array_targets)) + return {target: self.generate_arrays(target) for target in targets} def generate_time_mapper(funcs): diff --git a/devito/petsc/types/array.py b/devito/petsc/types/array.py index 049785662b..8b57aca44c 100644 --- a/devito/petsc/types/array.py +++ b/devito/petsc/types/array.py @@ -69,6 +69,10 @@ def dimensions(self): def target(self): return self._target + @property + def grid(self): + return self.target.grid + @property def coefficients(self): """Form of the coefficients of the function.""" diff --git a/devito/petsc/types/equation.py b/devito/petsc/types/equation.py index a5c5bc55ef..ad929c8ed5 100644 --- a/devito/petsc/types/equation.py +++ b/devito/petsc/types/equation.py @@ -6,33 +6,38 @@ class EssentialBC(Eq): """ - A special equation used to handle essential boundary conditions - in the PETSc solver. Until PetscSection + DMDA is supported, - we treat essential boundary conditions as trivial equations - in the solver, where we place 1.0 (scaled) on the diagonal of - the jacobian, zero symmetrically and move the boundary - data to the right-hand side. - - NOTE: When users define essential boundary conditions, they need to ensure that - the SubDomains do not overlap. Solver will still run but may see unexpected behaviour - at boundaries. This will be documented in the PETSc examples. + Represents an essential boundary condition for use with PETScSolve. + + Due to ongoing work on PetscSection and DMDA integration (WIP), + these conditions are imposed as trivial equations. The compiler + will automatically zero the corresponding rows/columns in the Jacobian + and lift the boundary terms into the residual RHS. + + Note: + - To define an essential boundary condition, use: + Eq(target, boundary_value, subdomain=...), + where `target` is the Function-like object passed to PETScSolve. + - SubDomains used for multiple EssentialBCs must not overlap. """ pass class ZeroRow(EssentialBC): """ - Equation used to zero the row of the Jacobian corresponding - to an essential BC. - This is only used interally by the compiler, not by users. + Equation used to zero all entries, except the diagonal, + of a row in the Jacobian. + + Note: + This is only used interally by the compiler, not by users. """ pass class ZeroColumn(EssentialBC): """ - Equation used to zero the column of the Jacobian corresponding - to an essential BC. - This is only used interally by the compiler, not by users. + Equation used to zero the column of the Jacobian. + + Note: + This is only used interally by the compiler, not by users. """ - pass \ No newline at end of file + pass diff --git a/devito/petsc/types/object.py b/devito/petsc/types/object.py index e62d125dc2..32892f2aa4 100644 --- a/devito/petsc/types/object.py +++ b/devito/petsc/types/object.py @@ -209,7 +209,6 @@ def __init__(self, name='subctx', pname='SubMatrixCtx', fields=None, _C_modifier = None - class PETScArrayObject(PetscMixin, ArrayObject): _data_alignment = False diff --git a/devito/petsc/types/types.py b/devito/petsc/types/types.py index 3897552f7b..a0a9cd2fe5 100644 --- a/devito/petsc/types/types.py +++ b/devito/petsc/types/types.py @@ -5,7 +5,7 @@ from devito.tools import Reconstructable, sympy_mutex, as_tuple from devito.tools.dtypes_lowering import dtype_mapper from devito.petsc.utils import petsc_variables -from devito.symbolics.extraction import separate_eqn, generate_targets +from devito.symbolics.extraction import separate_eqn, generate_targets, centre_stencil from devito.petsc.types.equation import EssentialBC, ZeroRow, ZeroColumn from devito.types.equation import Eq from devito.operations.solve import eval_time_derivatives @@ -186,7 +186,7 @@ def space_order(self): @property def targets(self): - return (self.target,) + return as_tuple(self.target) class MultipleFieldData(FieldData): @@ -229,22 +229,104 @@ def space_order(self): ) return space_orders.pop() - @property - def jacobian(self): - return self._jacobian - @property def targets(self): return self._targets - -class BaseJacobian: - def __init__(self, targets, time_mapper, arrays): - self.targets = as_tuple(targets) +class Jacobian: + def __init__(self, target, eqns, arrays, time_mapper): + self.target = target + self.eqns = eqns + self.arrays = arrays self.time_mapper = time_mapper + self._build_matvecs() + + @property + def matvecs(self): + return self._matvecs + + @property + def scdiag(self): + return self._scdiag + + def _build_matvecs(self): + matvecs = [ + e for eq in self.eqns for e in + self._build_matvec_eq(eq) + if e is not None + ] + matvecs = tuple(sorted(matvecs, key=lambda e: not isinstance(e, EssentialBC))) + + matvecs = self._scale_non_bcs(matvecs) + scdiag = self._compute_scdiag(matvecs) + matvecs = self._scale_bcs(matvecs, scdiag) + + self._matvecs = matvecs + self._scdiag = scdiag + + def _build_matvec_eq(self, eq, target=None, arrays=None): + target = target or self.target + arrays = arrays or self.arrays + + b, F_target, _, targets = separate_eqn(eq, target) + if F_target: + return self._make_matvec(eq, F_target, targets, arrays) + return (None,) + + def _make_matvec(self, eq, F_target, targets, arrays): + if isinstance(eq, EssentialBC): + # NOTE: Essential BCs are trivial equations in the solver. + # See `EssentialBC` for more details. + rhs = arrays['x'] + zero_row = ZeroRow(arrays['y'], rhs, subdomain=eq.subdomain) + zero_column = ZeroColumn(arrays['x'], 0., subdomain=eq.subdomain) + return (zero_row, zero_column) + else: + rhs = F_target.subs(targets_to_arrays(arrays['x'], targets)) + rhs = rhs.subs(self.time_mapper) + return as_tuple(Eq(arrays['y'], rhs, subdomain=eq.subdomain)) + + def _scale_non_bcs(self, matvecs, target=None): + target = target or self.target + vol = target.grid.symbolic_volume_cell + + return [ + m if isinstance(m, EssentialBC) else m._rebuild(rhs=m.rhs * vol) + for m in matvecs + ] + + def _compute_scdiag(self, matvecs, arrays=None): + """ + """ + arrays = arrays or self.arrays + + centres = { + centre_stencil(m.rhs, arrays['x'], as_coeff=True) + for m in matvecs if not isinstance(m, EssentialBC) + } + # add comments + return centres.pop() if len(centres) == 1 else 1.0 + + def _scale_bcs(self, matvecs, scdiag): + """ + Scale the essential BCs + """ + return [ + m._rebuild(rhs=m.rhs * scdiag) if isinstance(m, ZeroRow) else m + for m in matvecs + ] + + +class MixedJacobian(Jacobian): + def __init__(self, target_eqns, arrays, time_mapper): + """ + """ + self.targets = as_tuple(target_eqns.keys()) self.arrays = arrays + self.time_mapper = time_mapper self.submatrices = self._initialize_submatrices() + self._build_blocks(target_eqns) def _initialize_submatrices(self): """ @@ -260,11 +342,33 @@ def _initialize_submatrices(self): submatrices[target][key] = { 'matvecs': None, 'derivative_wrt': self.targets[j], - 'index': i * num_targets + j + 'index': i * num_targets + j, + 'scdiag': None } return submatrices + def _build_blocks(self, target_eqns): + for target, eqns in target_eqns.items(): + self._build_block(target, eqns) + + def _build_block(self, target, eqns): + arrays = self.arrays[target] + for submat, mtvs in self.submatrices[target].items(): + matvecs = [ + e for eq in eqns for e in + self._build_matvec_eq(eq, mtvs['derivative_wrt'], arrays) + ] + matvecs = [m for m in matvecs if m is not None] + matvecs = tuple(sorted(matvecs, key=lambda e: not isinstance(e, EssentialBC))) + + matvecs = self._scale_non_bcs(matvecs, target) + scdiag = self._compute_scdiag(matvecs, arrays) + matvecs = self._scale_bcs(matvecs, scdiag) + + if matvecs: + self.set_submatrix(target, submat, matvecs, scdiag) + @property def submatrix_keys(self): """ @@ -295,21 +399,21 @@ def submat_to_index(self): for key, value in submats.items() } - @property - def diagonal_submatrix_keys(self): + # CHECK/TEST THIS + def is_diagonal_submatrix(self, key): """ - Return a list of diagonal submatrix keys (e.g., ['J00', 'J11']). + Return True if the given key corresponds to a diagonal + submatrix (e.g., 'J00', 'J11'), else False. """ - keys = [] - for i, target in enumerate(self.targets): + for i, t in enumerate(self.targets): diag_key = f'J{i}{i}' - if diag_key in self.submatrices[target]: - keys.append(diag_key) - return keys + if key == diag_key and diag_key in self.submatrices[t]: + return True + return False - def set_submatrix(self, field, key, matvecs): + def set_submatrix(self, field, key, matvecs, scdiag): """ - Set a specific submatrix for a field. + Set the matrix-vector equations for a submatrix. Parameters ---------- @@ -321,7 +425,8 @@ def set_submatrix(self, field, key, matvecs): The matrix-vector equations forming the submatrix. """ if field in self.submatrices and key in self.submatrices[field]: - self.submatrices[field][key]["matvecs"] = matvecs + self.submatrices[field][key]['matvecs'] = matvecs + self.submatrices[field][key]['scdiag'] = scdiag else: raise KeyError(f'Invalid field ({field}) or submatrix key ({key})') @@ -331,112 +436,63 @@ def get_submatrix(self, field, key): """ return self.submatrices.get(field, {}).get(key, None) - def build_matvec_eq(self, eq, target, arrays): - b, F_target, _, targets = separate_eqn(eq, target) - if F_target: - return self.make_matvec(eq, F_target, targets, arrays) - return (None,) - - def make_matvec(self, eq, F_target, targets, arrays): - if isinstance(eq, EssentialBC): - # NOTE: Until PetscSection + DMDA is supported, we leave - # the essential BCs in the solver. - # Trivial equations for bc rows -> place 1.0 on diagonal (scaled) - # and zero symmetrically. - rhs = arrays['x'] - zero_row = ZeroRow(arrays['y'], rhs, subdomain=eq.subdomain) - zero_column = ZeroColumn(arrays['x'], 0.0, subdomain=eq.subdomain) - return (zero_row, zero_column) - else: - rhs = F_target.subs(targets_to_arrays(arrays['x'], targets)) - # rhs = rhs.subs(self.time_mapper) * self.cell_area - rhs = rhs = rhs.subs(self.time_mapper) - - return as_tuple(Eq(arrays['y'], rhs, subdomain=eq.subdomain)) + @property + def target_scaler_mapper(self): + """ + Return a mapping from each target to its diagonal submatrix's scaler. + """ + mapper = {} + for target, submats in self.submatrices.items(): + for key, value in submats.items(): + if self.is_diagonal_submatrix(key): + mapper[target] = value.get('scdiag') + break + return mapper def __repr__(self): + # TODO: edit return str(self.submatrices) -class Jacobian(BaseJacobian): +class Residual: + """ + """ + def __init__(self, target, eqns, arrays, time_mapper, scdiag): + self.target = target + self.eqns = eqns + self.arrays = arrays + self.time_mapper = time_mapper + self.scdiag = scdiag + self._build_equations() @property - def target(self): - return self.targets[0] + def formfuncs(self): + """ + """ + return self._formfuncs @property - def matvecs(self): - return self.submatrices[self.target]['J00']['matvecs'] - - # TODO: use same structure arrays for both jacobian and mixedjacobian - def build_block(self, eqns): - for submat, mtvs in self.submatrices[self.target].items(): - matvecs = [ - e for eq in eqns for e in - self.build_matvec_eq(eq, mtvs['derivative_wrt'], self.arrays) - ] - matvecs = [m for m in matvecs if m is not None] - matvecs = tuple(sorted(matvecs, key=lambda e: not isinstance(e, EssentialBC))) - - if matvecs: - self.set_submatrix(self.target, submat, matvecs) - - -class MixedJacobian(BaseJacobian): - - # TODO: use same structure arrays for both jacobian and mixedjacobian - def build_block(self, target, eqns): - for submat, mtvs in self.submatrices[target].items(): - matvecs = [ - e for eq in eqns for e in - self.build_matvec_eq(eq, mtvs['derivative_wrt'], self.arrays[target]) - ] - matvecs = [m for m in matvecs if m is not None] - matvecs = tuple(sorted(matvecs, key=lambda e: not isinstance(e, EssentialBC))) - - if matvecs: - self.set_submatrix(target, submat, matvecs) - - def build_blocks(self, target_eqns): - for target, eqns in target_eqns.items(): - self.build_block(target, eqns) - - -class BaseResidual: - def scale_essential_bcs(self, equations): + def formrhs(self): """ """ - return [ - eq._rebuild(rhs=self.scale * eq.rhs) if isinstance(eq, ZeroRow) else eq - for eq in equations - ] + return self._formrhs - -class Residual(BaseResidual): - """ - """ - - def __init__(self, target, time_mapper, arrays, scale): - self.target = target - self.time_mapper = time_mapper - self.arrays = arrays - self.scale = scale - self.formfuncs = [] - self.formrhs = [] - - def build_equations(self, eqns): + def _build_equations(self): """ """ - for eq in eqns: + funcs = [] + rhs = [] + + for eq in self.eqns: b, F_target, _, targets = separate_eqn(eq, self.target) - F_target = self.make_F_target(eq, F_target, targets) - b = self.make_b(eq, b) - self.formfuncs.extend(F_target) - self.formrhs.extend(b) + funcs.extend(self._make_F_target(eq, F_target, targets)) + # TODO: if b is zero then don't need a rhs vector+callback + rhs.extend(self._make_b(eq, b)) - self.formfuncs = self.scale_essential_bcs(self.formfuncs) + self._formfuncs = [self._scale_bcs(eq) for eq in funcs] + self._formrhs = rhs - def make_F_target(self, eq, F_target, targets): + def _make_F_target(self, eq, F_target, targets): arrays = self.arrays volume = self.target.grid.symbolic_volume_cell if isinstance(eq, EssentialBC): @@ -455,55 +511,64 @@ def make_F_target(self, eq, F_target, targets): rhs = rhs.subs(self.time_mapper) * volume return as_tuple(Eq(arrays['f'], rhs, subdomain=eq.subdomain)) - def make_b(self, eq, b): + def _make_b(self, eq, b): rhs = 0. if isinstance(eq, EssentialBC) else b.subs(self.time_mapper) rhs = rhs * self.target.grid.symbolic_volume_cell return as_tuple(Eq(self.arrays['b'], rhs, subdomain=eq.subdomain)) + def _scale_bcs(self, eq, scdiag=None): + """ + Scale ZeroRow equations using scdiag + """ + scdiag = scdiag or self.scdiag + return eq._rebuild(rhs=scdiag * eq.rhs) if isinstance(eq, ZeroRow) else eq + -class MixedResidual(BaseResidual): +class MixedResidual(Residual): """ """ - # TODO: change default and pass in correct scale - def __init__(self, targets, time_mapper, arrays, scale=1.0): - self.targets = as_tuple(targets) - self.time_mapper = time_mapper + def __init__(self, target_eqns, arrays, time_mapper, scdiag): + self.targets = as_tuple(target_eqns.keys()) self.arrays = arrays - self.scale = scale - self.formfuncs = [] + self.time_mapper = time_mapper + self.scdiag = scdiag + self._build_equations(target_eqns) - def build_equations(self, eqn_dict): + @property + def formrhs(self): + """ + """ + return None + + def _build_equations(self, target_eqns): all_formfuncs = [] - for target, eqns in eqn_dict.items(): + for target, eqns in target_eqns.items(): formfuncs = chain.from_iterable( - self.build_function_eq(eq, target) + self._build_function_eq(eq, target) for eq in as_tuple(eqns) ) + all_formfuncs.extend(formfuncs) - # scale, = self._diag_scale[arrays[target]['x']] - # fix this - # scale = 1.0 - all_formfuncs.extend(self.scale_essential_bcs(formfuncs)) - - self.formfuncs = tuple(sorted( + self._formfuncs = tuple(sorted( all_formfuncs, key=lambda e: not isinstance(e, EssentialBC) )) - - def build_function_eq(self, eq, target): + def _build_function_eq(self, eq, target): zeroed = eq.lhs - eq.rhs zeroed_eqn = Eq(eq.lhs - eq.rhs, 0) eval_zeroed_eqn = eval_time_derivatives(zeroed_eqn.lhs) + volume = target.grid.symbolic_volume_cell + mapper = {} for t in self.targets: target_funcs = set(generate_targets(Eq(eval_zeroed_eqn, 0), t)) mapper.update(targets_to_arrays(self.arrays[t]['x'], target_funcs)) if isinstance(eq, EssentialBC): - rhs = self.arrays[target]['x'] - eq.rhs + rhs = (self.arrays[target]['x'] - eq.rhs)*self.scdiag[target] zero_row = ZeroRow( self.arrays[target]['f'], rhs, subdomain=eq.subdomain ) @@ -513,16 +578,51 @@ def build_function_eq(self, eq, target): return (zero_row, zero_col) else: if isinstance(zeroed, (int, float)): - # rhs = zeroed * self.cell_area - rhs = zeroed + rhs = zeroed * volume else: rhs = zeroed.subs(mapper) - # rhs = rhs.subs(self.time_mapper)*self.cell_area - rhs = rhs.subs(self.time_mapper) + rhs = rhs.subs(self.time_mapper)*volume return as_tuple(Eq(self.arrays[target]['f'], rhs, subdomain=eq.subdomain)) +class InitialGuess: + """ + Enforce initial guess to satisfy essential BCs. + # TODO: Extend this to "coupled". + """ + def __init__(self, target, eqns, arrays): + self.target = target + self.eqns = as_tuple(eqns) + self.arrays = arrays + self._build_equations() + + @property + def equations(self): + """ + """ + return self._equations + + def _build_equations(self): + """ + Return a list of initial guess equations. + """ + self._equations = [ + eq for eq in + (self._make_initial_guess(e) for e in self.eqns) + if eq is not None + ] + + def _make_initial_guess(self, eq): + if isinstance(eq, EssentialBC): + assert eq.lhs == self.target + return Eq( + self.arrays['x'], eq.rhs, + subdomain=eq.subdomain + ) + else: + return None + def targets_to_arrays(array, targets): """ diff --git a/examples/petsc/Poisson/01_poisson.py b/examples/petsc/Poisson/01_poisson.py index 70f17d94c6..4981566af6 100644 --- a/examples/petsc/Poisson/01_poisson.py +++ b/examples/petsc/Poisson/01_poisson.py @@ -66,7 +66,6 @@ def analytical(x, y): Ly = np.float64(1.) n_values = list(range(13, 174, 10)) -n_values = [7] dx = np.array([Lx/(n-1) for n in n_values]) errors = [] @@ -112,7 +111,7 @@ def analytical(x, y): error = np.linalg.norm(diff) / np.linalg.norm(phi_analytical[1:-1, 1:-1]) errors.append(error) -# slope, _ = np.polyfit(np.log(dx), np.log(errors), 1) +slope, _ = np.polyfit(np.log(dx), np.log(errors), 1) -# assert slope > 1.9 -# assert slope < 2.1 +assert slope > 1.9 +assert slope < 2.1 diff --git a/examples/petsc/petsc_test.py b/examples/petsc/petsc_test.py index 5d93669d5f..12d245480e 100644 --- a/examples/petsc/petsc_test.py +++ b/examples/petsc/petsc_test.py @@ -20,7 +20,7 @@ v.data[:] = 5.0 -eq = Eq(v, u.laplace, subdomain=grid.interior) +eq = Eq(0., u.laplace, subdomain=grid.interior) petsc = PETScSolve([eq], u) diff --git a/examples/petsc/random/biharmonic/02_biharmonic.py b/examples/petsc/random/02_biharmonic.py similarity index 100% rename from examples/petsc/random/biharmonic/02_biharmonic.py rename to examples/petsc/random/02_biharmonic.py diff --git a/examples/petsc/random/biharmonic/biharmonic_matfree.c b/examples/petsc/random/biharmonic/biharmonic_matfree.c deleted file mode 100644 index 2fb432382f..0000000000 --- a/examples/petsc/random/biharmonic/biharmonic_matfree.c +++ /dev/null @@ -1,556 +0,0 @@ - -// # ref - https://github.com/bueler/p4pdes/blob/master/c/ch7/biharm.c - -static char help[] = -"Solve the linear biharmonic equation in 2D. Equation is\n" -" Lap^2 u = f\n" -"where Lap = - grad^2 is the positive Laplacian, equivalently\n" -" u_xxxx + 2 u_xxyy + u_yyyy = f(x,y)\n" -"Domain is unit square S = (0,1)^2. Boundary conditions are homogeneous\n" -"simply-supported: u = 0, Lap u = 0. The equation is rewritten as a\n" -"2x2 block system with SPD Laplacian blocks on the diagonal:\n" -" | Lap | 0 | | v | | f | \n" -" |-----|-----| |---| = |---| \n" -" | -I | Lap | | u | | 0 | \n" -"Includes manufactured, polynomial exact solution. The discretization is\n" -"structured-grid (DMDA) finite differences. Includes analytical Jacobian.\n" -"Recommended preconditioning combines fieldsplit:\n" -" -pc_type fieldsplit -pc_fieldsplit_type multiplicative|additive \n" -"with multigrid as the preconditioner for the diagonal blocks:\n" -" -fieldsplit_v_pc_type mg|gamg -fieldsplit_u_pc_type mg|gamg\n" -"(GMG requires setting levels and Galerkin coarsening.) One can also do\n" -"monolithic multigrid (-pc_type mg|gamg).\n\n"; - -#include - -typedef struct { - PetscReal v, u; -} Field; - -typedef struct { - PetscReal (*f)(PetscReal x, PetscReal y); // right-hand side -} BiharmCtx; - -struct JacobianCtx -{ - DM * subdms; - IS * fields; - Mat * submats; -} ; - -struct SubMatrixCtx -{ - IS * rows; - IS * cols; -} ; - -static PetscReal c(PetscReal x) { - return x*x*x * (1.0-x)*(1.0-x)*(1.0-x); -} - -static PetscReal ddc(PetscReal x) { - return 6.0 * x * (1.0-x) * (1.0 - 5.0 * x + 5.0 * x*x); -} - -static PetscReal d4c(PetscReal x) { - return - 72.0 * (1.0 - 5.0 * x + 5.0 * x*x); -} - -static PetscReal u_exact_fcn(PetscReal x, PetscReal y) { - return c(x) * c(y); -} - -static PetscReal lap_u_exact_fcn(PetscReal x, PetscReal y) { - return - ddc(x) * c(y) - c(x) * ddc(y); // Lap u = - grad^2 u -} - -static PetscReal f_fcn(PetscReal x, PetscReal y) { - return d4c(x) * c(y) + 2.0 * ddc(x) * ddc(y) + c(x) * d4c(y); // Lap^2 u = grad^4 u -} - -extern PetscErrorCode FormExactWLocal(DMDALocalInfo*, Field**, BiharmCtx*); -extern PetscErrorCode FormFunction(SNES snes, Vec X, Vec F, void* dummy); -extern PetscErrorCode J00_MatMult(Mat J, Vec X, Vec Y); -extern PetscErrorCode J10_MatMult(Mat J, Vec X, Vec Y); -extern PetscErrorCode J11_MatMult(Mat J, Vec X, Vec Y); -extern PetscErrorCode WholeMatMult(Mat J, Vec X, Vec Y); -PetscErrorCode MatCreateSubMatrices0(Mat J, PetscInt nfields, IS * irow, IS * icol, MatReuse scall, Mat * * submats); -extern PetscErrorCode PopulateMatContext(struct JacobianCtx * jctx, DM * subdms, IS * fields); - -int main(int argc,char **argv) { - DM da; - SNES snes; - Vec w, w_initial, w_exact; - BiharmCtx user; - Field **aW; - PetscReal normv, normu, errv, erru; - DMDALocalInfo info; - IS *fields; - DM *subdms; - PetscInt nfields; - - struct JacobianCtx jctx0; - Mat J; - - PetscCall(PetscInitialize(&argc,&argv,NULL,help)); - - user.f = &f_fcn; - PetscCall(DMDACreate2d(PETSC_COMM_WORLD, - DM_BOUNDARY_NONE, DM_BOUNDARY_NONE, DMDA_STENCIL_STAR, - 6,6,PETSC_DECIDE,PETSC_DECIDE, - 2,1, // degrees of freedom, stencil width - NULL,NULL,&da)); - PetscCall(DMSetApplicationContext(da,&user)); - PetscCall(DMSetFromOptions(da)); - PetscCall(DMSetUp(da)); // this must be called BEFORE SetUniformCoordinates - PetscCall(DMSetMatType(da, MATSHELL)); - PetscCall(DMDASetUniformCoordinates(da,0.0,1.0,0.0,1.0,-1.0,-1.0)); - PetscCall(DMDASetFieldName(da,0,"v")); - PetscCall(DMDASetFieldName(da,1,"u")); - PetscCall(DMCreateMatrix(da,&J)); - - PetscCall(SNESCreate(PETSC_COMM_WORLD,&snes)); - PetscCall(SNESSetDM(snes,da)); - PetscCall(SNESSetFunction(snes,NULL,FormFunction,NULL)); - PetscCall(SNESSetType(snes,SNESKSPONLY)); - PetscCall(SNESSetFromOptions(snes)); - - PetscCall(SNESSetJacobian(snes,J,J,MatMFFDComputeJacobian,NULL)); - PetscCall(MatShellSetOperation(J,MATOP_MULT,(void (*)(void))WholeMatMult)); - - PetscCall(MatSetDM(J,da)); - PetscCall(DMCreateFieldDecomposition(da,&(nfields),NULL,&fields,&subdms)); - PetscCall(PopulateMatContext(&(jctx0),subdms,fields)); - PetscCall(MatShellSetContext(J,&(jctx0))); - PetscCall(MatCreateSubMatrices0(J,nfields,fields,fields,MAT_INITIAL_MATRIX,&(jctx0.submats))); - - PetscCall(DMGetGlobalVector(da,&w_initial)); - PetscCall(VecSet(w_initial,0.0)); - PetscCall(SNESSolve(snes,NULL,w_initial)); - // PetscCall(VecView(w_initial,PETSC_VIEWER_STDOUT_WORLD)); - PetscCall(DMRestoreGlobalVector(da,&w_initial)); - PetscCall(DMDestroy(&da)); - - PetscCall(SNESGetSolution(snes,&w)); - PetscCall(SNESGetDM(snes,&da)); - PetscCall(DMDAGetLocalInfo(da,&info)); - - PetscCall(DMCreateGlobalVector(da,&w_exact)); - PetscCall(DMDAVecGetArray(da,w_exact,&aW)); - PetscCall(FormExactWLocal(&info,aW,&user)); - PetscCall(DMDAVecRestoreArray(da,w_exact,&aW)); - PetscCall(VecStrideNorm(w_exact,0,NORM_INFINITY,&normv)); - PetscCall(VecStrideNorm(w_exact,1,NORM_INFINITY,&normu)); - PetscCall(VecAXPY(w,-1.0,w_exact)); - PetscCall(VecStrideNorm(w,0,NORM_INFINITY,&errv)); - PetscCall(VecStrideNorm(w,1,NORM_INFINITY,&erru)); - PetscCall(PetscPrintf(PETSC_COMM_WORLD, - "done on %d x %d grid ...\n" - " errors |v-vex|_inf/|vex|_inf = %.5e, |u-uex|_inf/|uex|_inf = %.5e\n", - info.mx,info.my,errv/normv,erru/normu)); - - - PetscCall(ISDestroy(&(fields[0]))); - PetscCall(ISDestroy(&(fields[1]))); - PetscCall(PetscFree(fields)); - PetscCall(DMDestroy(&(subdms[0]))); - PetscCall(DMDestroy(&(subdms[1]))); - PetscCall(PetscFree(subdms)); - PetscCall(VecDestroy(&w_exact)); - PetscCall(MatDestroy(&J)); - PetscCall(SNESDestroy(&snes)); - PetscCall(PetscFinalize()); - return 0; -} - -PetscErrorCode FormExactWLocal(DMDALocalInfo *info, Field **aW, BiharmCtx *user) { - PetscInt i, j; - PetscReal xymin[2], xymax[2], hx, hy, x, y; - PetscCall(DMGetBoundingBox(info->da,xymin,xymax)); - hx = (xymax[0] - xymin[0]) / (info->mx - 1); - hy = (xymax[1] - xymin[1]) / (info->my - 1); - for (j = info->ys; j < info->ys + info->ym; j++) { - y = j * hy; - for (i = info->xs; i < info->xs + info->xm; i++) { - x = i * hx; - aW[j][i].u = u_exact_fcn(x,y); - aW[j][i].v = lap_u_exact_fcn(x,y); - } - } - return 0; -} - - -PetscErrorCode FormFunction(SNES snes, Vec X, Vec F, void * dummy) -{ - Vec xlocal, flocal; - DMDALocalInfo info; - DM da; - PetscScalar *x_vec, *f_vec; - - BiharmCtx *user; - - PetscCall(SNESGetDM(snes,&da)); - - PetscCall(DMGetApplicationContext(da,&user)); - - PetscCall(DMDAGetLocalInfo(da,&info)); - PetscCall(DMGetLocalVector(da,&xlocal)); - PetscCall(DMGetLocalVector(da,&flocal)); - - PetscCall(DMGlobalToLocalBegin(da,X,INSERT_VALUES,xlocal)); - PetscCall(DMGlobalToLocalEnd(da,X,INSERT_VALUES,xlocal)); - - PetscCall(VecGetArray(xlocal,&x_vec)); - PetscCall(VecGetArray(flocal,&f_vec)); - - Field (*xx)[info.gxm] = (Field (*)[info.gxm]) x_vec; - Field (*ff)[info.gxm] = (Field (*)[info.gxm]) f_vec; - - PetscInt i, j; - PetscReal xymin[2], xymax[2], hx, hy, darea, scx, scy, scdiag, x, y, - ve, vw, vn, vs, ue, uw, un, us; - - hx = 1. / (info.mx - 1); - hy = 1. / (info.my - 1); - - darea = hx * hy; // multiply FD equations by this - - scx = hy / hx; - scy = hx / hy; - scdiag = 2.0 * (scx + scy); // diagonal scaling - for (j = info.ys; j < info.ys + info.ym; j++) { - y = xymin[1] + j * hy; - for (i = info.xs; i < info.xs + info.xm; i++) { - x = xymin[0] + i * hx; - if (i==0 || i==info.mx-1 || j==0 || j==info.my-1) { - ff[j][i].v = scdiag * xx[j][i].v; - ff[j][i].u = scdiag * xx[j][i].u; - } else { - ve = xx[j][i+1].v; - vw = xx[j][i-1].v; - vn = xx[j+1][i].v; - vs = xx[j-1][i].v; - ff[j][i].v = scdiag * xx[j][i].v - scx * (vw + ve) - scy * (vs + vn) - - darea * (*(user->f))(x,y); - ue = xx[j][i+1].u; - uw = xx[j][i-1].u; - un = xx[j+1][i].u; - us = xx[j-1][i].u; - ff[j][i].u = - darea * xx[j][i].v - + scdiag * xx[j][i].u - scx * (uw + ue) - scy * (us + un); - } - } - } - - PetscCall(VecRestoreArray(xlocal,&x_vec)); - PetscCall(VecRestoreArray(flocal,&f_vec)); - - PetscCall(DMLocalToGlobalBegin(da,flocal,INSERT_VALUES,F)); - PetscCall(DMLocalToGlobalEnd(da,flocal,INSERT_VALUES,F)); - PetscCall(DMRestoreLocalVector(da,&xlocal)); - PetscCall(DMRestoreLocalVector(da,&flocal)); - - return 0; -} - - -PetscErrorCode J00_MatMult(Mat J, Vec X, Vec Y) -{ - PetscFunctionBeginUser; - - DM dm0; - DMDALocalInfo info; - Vec xloc; - Vec yloc; - - BiharmCtx * ctx0; - PetscScalar * x_v_vec; - PetscScalar * y_v_vec; - - PetscCall(MatGetDM(J,&(dm0))); - PetscCall(DMGetApplicationContext(dm0,&(ctx0))); - PetscCall(DMGetLocalVector(dm0,&(xloc))); - PetscCall(DMGlobalToLocalBegin(dm0,X,INSERT_VALUES,xloc)); - PetscCall(DMGlobalToLocalEnd(dm0,X,INSERT_VALUES,xloc)); - PetscCall(DMGetLocalVector(dm0,&(yloc))); - PetscCall(VecSet(yloc,0.0)); - PetscCall(VecGetArray(yloc,&y_v_vec)); - PetscCall(VecGetArray(xloc,&x_v_vec)); - PetscCall(DMDAGetLocalInfo(dm0,&(info))); - - PetscScalar (* x_v)[info.gxm] = (PetscScalar (*)[info.gxm]) x_v_vec; - PetscScalar (* y_v)[info.gxm] = (PetscScalar (*)[info.gxm]) y_v_vec; - - PetscInt i, j; - PetscReal xymin[2], xymax[2], hx, hy, darea, scx, scy, scdiag, x, y, - ve, vw, vn, vs, ue, uw, un, us; - - hx = 1./ (info.mx - 1); - hy = 1./ (info.my - 1); - darea = hx * hy; // multiply FD equations by this - scx = hy / hx; - scy = hx / hy; - scdiag = 2.0 * (scx + scy); // diagonal scaling - for (j = info.ys; j < info.ys + info.ym; j++) { - y = xymin[1] + j * hy; - for (i = info.xs; i < info.xs + info.xm; i++) { - x = xymin[0] + i * hx; - if (i==0 || i==info.mx-1 || j==0 || j==info.my-1) { - y_v[j][i] = scdiag * x_v[j][i]; - } else { - ve = x_v[j][i+1]; - vw = x_v[j][i-1]; - vn = x_v[j+1][i]; - vs = x_v[j-1][i]; - y_v[j][i] = scdiag * x_v[j][i] - scx * (vw + ve) - scy * (vs + vn); - - } - } - } - - PetscCall(VecRestoreArray(yloc,&y_v_vec)); - PetscCall(VecRestoreArray(xloc,&x_v_vec)); - PetscCall(DMLocalToGlobalBegin(dm0,yloc,ADD_VALUES,Y)); - PetscCall(DMLocalToGlobalEnd(dm0,yloc,ADD_VALUES,Y)); - PetscCall(DMRestoreLocalVector(dm0,&(xloc))); - PetscCall(DMRestoreLocalVector(dm0,&(yloc))); - - PetscFunctionReturn(0); -} - -PetscErrorCode J10_MatMult(Mat J, Vec X, Vec Y) -{ - PetscFunctionBeginUser; - - DM dm0; - DMDALocalInfo info; - Vec xloc; - Vec yloc; - - BiharmCtx * ctx0; - PetscScalar * x_v_vec; - PetscScalar * y_v_vec; - - PetscCall(MatGetDM(J,&(dm0))); - PetscCall(DMGetApplicationContext(dm0,&(ctx0))); - PetscCall(DMGetLocalVector(dm0,&(xloc))); - PetscCall(DMGlobalToLocalBegin(dm0,X,INSERT_VALUES,xloc)); - PetscCall(DMGlobalToLocalEnd(dm0,X,INSERT_VALUES,xloc)); - PetscCall(DMGetLocalVector(dm0,&(yloc))); - PetscCall(VecSet(yloc,0.0)); - PetscCall(VecGetArray(yloc,&y_v_vec)); - PetscCall(VecGetArray(xloc,&x_v_vec)); - PetscCall(DMDAGetLocalInfo(dm0,&(info))); - - PetscScalar (* x_v)[info.gxm] = (PetscScalar (*)[info.gxm]) x_v_vec; - PetscScalar (* y_v)[info.gxm] = (PetscScalar (*)[info.gxm]) y_v_vec; - - PetscInt i, j; - PetscReal xymin[2], xymax[2], hx, hy, darea, scx, scy, scdiag, x, y, - ve, vw, vn, vs, ue, uw, un, us; - - hx = 1. / (info.mx - 1); - hy = 1. / (info.my - 1); - - darea = hx * hy; // multiply FD equations by this - scx = hy / hx; - scy = hx / hy; - scdiag = 2.0 * (scx + scy); // diagonal scaling - for (j = info.ys; j < info.ys + info.ym; j++) { - y = xymin[1] + j * hy; - for (i = info.xs; i < info.xs + info.xm; i++) { - x = xymin[0] + i * hx; - if (i==0 || i==info.mx-1 || j==0 || j==info.my-1) { - y_v[j][i] = 0.0; - } else { - y_v[j][i] = -darea * x_v[j][i]; - - } - } - } - - PetscCall(VecRestoreArray(yloc,&y_v_vec)); - PetscCall(VecRestoreArray(xloc,&x_v_vec)); - PetscCall(DMLocalToGlobalBegin(dm0,yloc,ADD_VALUES,Y)); - PetscCall(DMLocalToGlobalEnd(dm0,yloc,ADD_VALUES,Y)); - PetscCall(DMRestoreLocalVector(dm0,&(xloc))); - PetscCall(DMRestoreLocalVector(dm0,&(yloc))); - - PetscFunctionReturn(0); -} - - -PetscErrorCode J11_MatMult(Mat J, Vec X, Vec Y) -{ - PetscFunctionBeginUser; - - DM dm0; - DMDALocalInfo info; - Vec xloc; - Vec yloc; - - BiharmCtx * ctx0; - PetscScalar * x_v_vec; - PetscScalar * y_v_vec; - - PetscCall(MatGetDM(J,&(dm0))); - PetscCall(DMGetApplicationContext(dm0,&(ctx0))); - PetscCall(DMGetLocalVector(dm0,&(xloc))); - PetscCall(DMGlobalToLocalBegin(dm0,X,INSERT_VALUES,xloc)); - PetscCall(DMGlobalToLocalEnd(dm0,X,INSERT_VALUES,xloc)); - PetscCall(DMGetLocalVector(dm0,&(yloc))); - PetscCall(VecSet(yloc,0.0)); - PetscCall(VecGetArray(yloc,&y_v_vec)); - PetscCall(VecGetArray(xloc,&x_v_vec)); - PetscCall(DMDAGetLocalInfo(dm0,&(info))); - - PetscScalar (* x_v)[info.gxm] = (PetscScalar (*)[info.gxm]) x_v_vec; - PetscScalar (* y_v)[info.gxm] = (PetscScalar (*)[info.gxm]) y_v_vec; - - PetscInt i, j; - PetscReal xymin[2], xymax[2], hx, hy, darea, scx, scy, scdiag, x, y, - ve, vw, vn, vs, ue, uw, un, us; -// PetscCall(DMGetBoundingBox(info.da,xymin,xymax)); - hx = 1. / (info.mx - 1); - hy = 1. / (info.my - 1); - - darea = hx * hy; - scx = hy / hx; - scy = hx / hy; - scdiag = 2.0 * (scx + scy); - for (j = info.ys; j < info.ys + info.ym; j++) { - y = xymin[1] + j * hy; - for (i = info.xs; i < info.xs + info.xm; i++) { - x = xymin[0] + i * hx; - if (i==0 || i==info.mx-1 || j==0 || j==info.my-1) { - y_v[j][i] = scdiag * x_v[j][i]; - } else { - ve = x_v[j][i+1]; - vw = x_v[j][i-1]; - vn = x_v[j+1][i]; - vs = x_v[j-1][i]; - y_v[j][i] = scdiag * x_v[j][i] - scx * (vw + ve) - scy * (vs + vn); - - } - } - } - - PetscCall(VecRestoreArray(yloc,&y_v_vec)); - PetscCall(VecRestoreArray(xloc,&x_v_vec)); - PetscCall(DMLocalToGlobalBegin(dm0,yloc,ADD_VALUES,Y)); - PetscCall(DMLocalToGlobalEnd(dm0,yloc,ADD_VALUES,Y)); - PetscCall(DMRestoreLocalVector(dm0,&(xloc))); - PetscCall(DMRestoreLocalVector(dm0,&(yloc))); - - PetscFunctionReturn(0); -} - -PetscErrorCode WholeMatMult(Mat J, Vec X, Vec Y) -{ - Vec J00X; - Vec J00Y; - Vec J10X; - Vec J10Y; - Vec J11X; - Vec J11Y; - - struct SubMatrixCtx * J00ctx; - struct SubMatrixCtx * J10ctx; - struct SubMatrixCtx * J11ctx; - struct JacobianCtx * jctx; - - PetscFunctionBeginUser; - - PetscCall(MatShellGetContext(J,&(jctx))); - - PetscCall(VecSet(Y,0.0)); - - Mat J00 = jctx->submats[0]; - PetscCall(MatShellGetContext(J00,&(J00ctx))); - PetscCall(VecGetSubVector(X,*(J00ctx->cols),&(J00X))); - PetscCall(VecGetSubVector(Y,*(J00ctx->rows),&(J00Y))); - PetscCall(MatMult(J00,J00X,J00Y)); - PetscCall(VecRestoreSubVector(X,*(J00ctx->cols),&(J00X))); - PetscCall(VecRestoreSubVector(Y,*(J00ctx->rows),&(J00Y))); - - Mat J10 = jctx->submats[2]; - PetscCall(MatShellGetContext(J10,&(J10ctx))); - PetscCall(VecGetSubVector(X,*(J10ctx->cols),&(J10X))); - PetscCall(VecGetSubVector(Y,*(J10ctx->rows),&(J10Y))); - PetscCall(MatMult(J10,J10X,J10Y)); - PetscCall(VecRestoreSubVector(X,*(J10ctx->cols),&(J10X))); - PetscCall(VecRestoreSubVector(Y,*(J10ctx->rows),&(J10Y))); - - Mat J11 = jctx->submats[3]; - PetscCall(MatShellGetContext(J11,&(J11ctx))); - PetscCall(VecGetSubVector(X,*(J11ctx->cols),&(J11X))); - PetscCall(VecGetSubVector(Y,*(J11ctx->rows),&(J11Y))); - PetscCall(MatMult(J11,J11X,J11Y)); - PetscCall(VecRestoreSubVector(X,*(J11ctx->cols),&(J11X))); - PetscCall(VecRestoreSubVector(Y,*(J11ctx->rows),&(J11Y))); - - - PetscFunctionReturn(0); -} - - -PetscErrorCode MatCreateSubMatrices0(Mat J, PetscInt nfields, IS * irow, IS * icol, MatReuse scall, Mat * * submats) -{ - PetscFunctionBeginUser; - - PetscInt M; - PetscInt N; - Mat block; - DM dm0; - PetscInt dof; - - struct UserCtx0 * ctx0; - struct JacobianCtx * jctx; - struct SubMatrixCtx * subctx; - - PetscCall(MatShellGetContext(J,&(jctx))); - DM * subdms = jctx->subdms; - - PetscInt nsubmats = nfields*nfields; - PetscCall(PetscCalloc1(nsubmats,submats)); - PetscCall(MatGetDM(J,&(dm0))); - PetscCall(DMGetApplicationContext(dm0,&(ctx0))); - PetscCall(DMDAGetInfo(dm0,NULL,&(M),&(N),NULL,NULL,NULL,NULL,&(dof),NULL,NULL,NULL,NULL,NULL)); - PetscInt subblockrows = M*N; - PetscInt subblockcols = M*N; - Mat * submat_arr = *submats; - - for (int i = 0; i <= nsubmats - 1; i += 1) - { - PetscCall(MatCreate(PETSC_COMM_WORLD,&(block))); - PetscCall(MatSetSizes(block,PETSC_DECIDE,PETSC_DECIDE,subblockrows,subblockcols)); - PetscCall(MatSetType(block,MATSHELL)); - PetscCall(PetscMalloc1(1,&(subctx))); - PetscInt rowidx = i / dof; - PetscInt colidx = (i)%(dof); - subctx->rows = &(irow[rowidx]); - subctx->cols = &(icol[colidx]); - PetscCall(DMSetApplicationContext(subdms[rowidx],ctx0)); - PetscCall(MatSetDM(block,subdms[rowidx])); - PetscCall(MatShellSetContext(block,subctx)); - PetscCall(MatSetUp(block)); - submat_arr[i] = block; - } - PetscCall(MatShellSetOperation(submat_arr[0],MATOP_MULT,(void (*)(void))J00_MatMult)); - PetscCall(MatShellSetOperation(submat_arr[2],MATOP_MULT,(void (*)(void))J10_MatMult)); - PetscCall(MatShellSetOperation(submat_arr[3],MATOP_MULT,(void (*)(void))J11_MatMult)); - - PetscFunctionReturn(0); -} - - -PetscErrorCode PopulateMatContext(struct JacobianCtx * jctx, DM * subdms, IS * fields) -{ - PetscFunctionBeginUser; - - jctx->subdms = subdms; - jctx->fields = fields; - - PetscFunctionReturn(0); -} diff --git a/examples/petsc/random/biharmonic/biharmonic_matfree_nonscaled.c b/examples/petsc/random/biharmonic/biharmonic_matfree_nonscaled.c deleted file mode 100644 index 04f2a02927..0000000000 --- a/examples/petsc/random/biharmonic/biharmonic_matfree_nonscaled.c +++ /dev/null @@ -1,553 +0,0 @@ - -// # ref - https://github.com/bueler/p4pdes/blob/master/c/ch7/biharm.c - -static char help[] = -"Solve the linear biharmonic equation in 2D. Equation is\n" -" Lap^2 u = f\n" -"where Lap = - grad^2 is the positive Laplacian, equivalently\n" -" u_xxxx + 2 u_xxyy + u_yyyy = f(x,y)\n" -"Domain is unit square S = (0,1)^2. Boundary conditions are homogeneous\n" -"simply-supported: u = 0, Lap u = 0. The equation is rewritten as a\n" -"2x2 block system with SPD Laplacian blocks on the diagonal:\n" -" | Lap | 0 | | v | | f | \n" -" |-----|-----| |---| = |---| \n" -" | -I | Lap | | u | | 0 | \n" -"Includes manufactured, polynomial exact solution. The discretization is\n" -"structured-grid (DMDA) finite differences. Includes analytical Jacobian.\n" -"Recommended preconditioning combines fieldsplit:\n" -" -pc_type fieldsplit -pc_fieldsplit_type multiplicative|additive \n" -"with multigrid as the preconditioner for the diagonal blocks:\n" -" -fieldsplit_v_pc_type mg|gamg -fieldsplit_u_pc_type mg|gamg\n" -"(GMG requires setting levels and Galerkin coarsening.) One can also do\n" -"monolithic multigrid (-pc_type mg|gamg).\n\n"; - -#include - -typedef struct { - PetscReal v, u; -} Field; - -typedef struct { - PetscReal (*f)(PetscReal x, PetscReal y); // right-hand side -} BiharmCtx; - -struct JacobianCtx -{ - DM * subdms; - IS * fields; - Mat * submats; -} ; - -struct SubMatrixCtx -{ - IS * rows; - IS * cols; -} ; - -static PetscReal c(PetscReal x) { - return x*x*x * (1.0-x)*(1.0-x)*(1.0-x); -} - -static PetscReal ddc(PetscReal x) { - return 6.0 * x * (1.0-x) * (1.0 - 5.0 * x + 5.0 * x*x); -} - -static PetscReal d4c(PetscReal x) { - return - 72.0 * (1.0 - 5.0 * x + 5.0 * x*x); -} - -static PetscReal u_exact_fcn(PetscReal x, PetscReal y) { - return c(x) * c(y); -} - -static PetscReal lap_u_exact_fcn(PetscReal x, PetscReal y) { - return - ddc(x) * c(y) - c(x) * ddc(y); // Lap u = - grad^2 u -} - -static PetscReal f_fcn(PetscReal x, PetscReal y) { - return d4c(x) * c(y) + 2.0 * ddc(x) * ddc(y) + c(x) * d4c(y); // Lap^2 u = grad^4 u -} - -extern PetscErrorCode FormExactWLocal(DMDALocalInfo*, Field**, BiharmCtx*); -extern PetscErrorCode FormFunction(SNES snes, Vec X, Vec F, void* dummy); -extern PetscErrorCode J00_MatMult(Mat J, Vec X, Vec Y); -extern PetscErrorCode J10_MatMult(Mat J, Vec X, Vec Y); -extern PetscErrorCode J11_MatMult(Mat J, Vec X, Vec Y); -extern PetscErrorCode WholeMatMult(Mat J, Vec X, Vec Y); -PetscErrorCode MatCreateSubMatrices0(Mat J, PetscInt nfields, IS * irow, IS * icol, MatReuse scall, Mat * * submats); -extern PetscErrorCode PopulateMatContext(struct JacobianCtx * jctx, DM * subdms, IS * fields); - -int main(int argc,char **argv) { - DM da; - SNES snes; - Vec w, w_initial, w_exact; - BiharmCtx user; - Field **aW; - PetscReal normv, normu, errv, erru; - DMDALocalInfo info; - IS *fields; - DM *subdms; - PetscInt nfields; - - struct JacobianCtx jctx0; - Mat J; - - PetscCall(PetscInitialize(&argc,&argv,NULL,help)); - - user.f = &f_fcn; - PetscCall(DMDACreate2d(PETSC_COMM_WORLD, - DM_BOUNDARY_NONE, DM_BOUNDARY_NONE, DMDA_STENCIL_STAR, - 33,33,PETSC_DECIDE,PETSC_DECIDE, - 2,1, // degrees of freedom, stencil width - NULL,NULL,&da)); - PetscCall(DMSetApplicationContext(da,&user)); - PetscCall(DMSetFromOptions(da)); - PetscCall(DMSetUp(da)); // this must be called BEFORE SetUniformCoordinates - PetscCall(DMSetMatType(da, MATSHELL)); - PetscCall(DMDASetUniformCoordinates(da,0.0,1.0,0.0,1.0,-1.0,-1.0)); - PetscCall(DMDASetFieldName(da,0,"v")); - PetscCall(DMDASetFieldName(da,1,"u")); - PetscCall(DMCreateMatrix(da,&J)); - - PetscCall(SNESCreate(PETSC_COMM_WORLD,&snes)); - PetscCall(SNESSetDM(snes,da)); - PetscCall(SNESSetFunction(snes,NULL,FormFunction,NULL)); - PetscCall(SNESSetType(snes,SNESKSPONLY)); - PetscCall(SNESSetFromOptions(snes)); - - PetscCall(SNESSetJacobian(snes,J,J,MatMFFDComputeJacobian,NULL)); - PetscCall(MatShellSetOperation(J,MATOP_MULT,(void (*)(void))WholeMatMult)); - - PetscCall(MatSetDM(J,da)); - PetscCall(DMCreateFieldDecomposition(da,&(nfields),NULL,&fields,&subdms)); - PetscCall(PopulateMatContext(&(jctx0),subdms,fields)); - PetscCall(MatShellSetContext(J,&(jctx0))); - PetscCall(MatCreateSubMatrices0(J,nfields,fields,fields,MAT_INITIAL_MATRIX,&(jctx0.submats))); - - PetscCall(DMGetGlobalVector(da,&w_initial)); - PetscCall(VecSet(w_initial,0.0)); - PetscCall(SNESSolve(snes,NULL,w_initial)); - // PetscCall(VecView(w_initial,PETSC_VIEWER_STDOUT_WORLD)); - PetscCall(DMRestoreGlobalVector(da,&w_initial)); - PetscCall(DMDestroy(&da)); - - PetscCall(SNESGetSolution(snes,&w)); - PetscCall(SNESGetDM(snes,&da)); - PetscCall(DMDAGetLocalInfo(da,&info)); - - PetscCall(DMCreateGlobalVector(da,&w_exact)); - PetscCall(DMDAVecGetArray(da,w_exact,&aW)); - PetscCall(FormExactWLocal(&info,aW,&user)); - - PetscCall(DMDAVecRestoreArray(da,w_exact,&aW)); - // PetscCall(VecView(w_exact,PETSC_VIEWER_STDOUT_WORLD)); - PetscCall(VecStrideNorm(w_exact,0,NORM_INFINITY,&normv)); - PetscCall(VecStrideNorm(w_exact,1,NORM_INFINITY,&normu)); - PetscCall(VecAXPY(w,-1.0,w_exact)); - PetscCall(VecStrideNorm(w,0,NORM_INFINITY,&errv)); - PetscCall(VecStrideNorm(w,1,NORM_INFINITY,&erru)); - PetscCall(PetscPrintf(PETSC_COMM_WORLD, - "done on %d x %d grid ...\n" - " errors |v-vex|_inf/|vex|_inf = %.5e, |u-uex|_inf/|uex|_inf = %.5e\n", - info.mx,info.my,errv/normv,erru/normu)); - - - PetscCall(ISDestroy(&(fields[0]))); - PetscCall(ISDestroy(&(fields[1]))); - PetscCall(PetscFree(fields)); - PetscCall(DMDestroy(&(subdms[0]))); - PetscCall(DMDestroy(&(subdms[1]))); - PetscCall(PetscFree(subdms)); - PetscCall(VecDestroy(&w_exact)); - PetscCall(MatDestroy(&J)); - PetscCall(SNESDestroy(&snes)); - PetscCall(PetscFinalize()); - return 0; -} - -PetscErrorCode FormExactWLocal(DMDALocalInfo *info, Field **aW, BiharmCtx *user) { - PetscInt i, j; - PetscReal xymin[2], xymax[2], hx, hy, x, y; - PetscCall(DMGetBoundingBox(info->da,xymin,xymax)); - hx = (xymax[0] - xymin[0]) / (info->mx - 1); - hy = (xymax[1] - xymin[1]) / (info->my - 1); - for (j = info->ys; j < info->ys + info->ym; j++) { - y = j * hy; - for (i = info->xs; i < info->xs + info->xm; i++) { - x = i * hx; - aW[j][i].u = u_exact_fcn(x,y); - aW[j][i].v = lap_u_exact_fcn(x,y); - } - } - return 0; -} - - -PetscErrorCode FormFunction(SNES snes, Vec X, Vec F, void * dummy) -{ - Vec xlocal, flocal; - DMDALocalInfo info; - DM da; - PetscScalar *x_vec, *f_vec; - - BiharmCtx *user; - - PetscCall(SNESGetDM(snes,&da)); - - PetscCall(DMGetApplicationContext(da,&user)); - - PetscCall(DMDAGetLocalInfo(da,&info)); - PetscCall(DMGetLocalVector(da,&xlocal)); - PetscCall(DMGetLocalVector(da,&flocal)); - - PetscCall(DMGlobalToLocalBegin(da,X,INSERT_VALUES,xlocal)); - PetscCall(DMGlobalToLocalEnd(da,X,INSERT_VALUES,xlocal)); - - PetscCall(VecGetArray(xlocal,&x_vec)); - PetscCall(VecGetArray(flocal,&f_vec)); - - Field (*xx)[info.gxm] = (Field (*)[info.gxm]) x_vec; - Field (*ff)[info.gxm] = (Field (*)[info.gxm]) f_vec; - - PetscInt i, j; - PetscReal xymin[2], xymax[2], hx, hy, darea, scx, scy, scdiag, x, y, - ve, vw, vn, vs, ue, uw, un, us; - - hx = 1. / (info.mx - 1); - hy = 1. / (info.my - 1); - - darea = hx * hy; // multiply FD equations by this - - scx = 1. / (hx*hx); - scy = 1. / (hy*hy); - scdiag = 2.0 * (scx + scy); // diagonal scaling - for (j = info.ys; j < info.ys + info.ym; j++) { - y = xymin[1] + j * hy; - for (i = info.xs; i < info.xs + info.xm; i++) { - x = xymin[0] + i * hx; - if (i==0 || i==info.mx-1 || j==0 || j==info.my-1) { - ff[j][i].v = xx[j][i].v; - ff[j][i].u = xx[j][i].u; - } else { - ve = xx[j][i+1].v; - vw = xx[j][i-1].v; - vn = xx[j+1][i].v; - vs = xx[j-1][i].v; - ff[j][i].v = scdiag * xx[j][i].v - scx * (vw + ve) - scy * (vs + vn) - - (*(user->f))(x,y); - ue = xx[j][i+1].u; - uw = xx[j][i-1].u; - un = xx[j+1][i].u; - us = xx[j-1][i].u; - ff[j][i].u = -xx[j][i].v - + scdiag * xx[j][i].u - scx * (uw + ue) - scy * (us + un); - } - } - } - - PetscCall(VecRestoreArray(xlocal,&x_vec)); - PetscCall(VecRestoreArray(flocal,&f_vec)); - - PetscCall(DMLocalToGlobalBegin(da,flocal,INSERT_VALUES,F)); - PetscCall(DMLocalToGlobalEnd(da,flocal,INSERT_VALUES,F)); - PetscCall(DMRestoreLocalVector(da,&xlocal)); - PetscCall(DMRestoreLocalVector(da,&flocal)); - - return 0; -} - - -PetscErrorCode J00_MatMult(Mat J, Vec X, Vec Y) -{ - PetscFunctionBeginUser; - - DM dm0; - DMDALocalInfo info; - Vec xloc; - Vec yloc; - - BiharmCtx * ctx0; - PetscScalar * x_v_vec; - PetscScalar * y_v_vec; - - PetscCall(MatGetDM(J,&(dm0))); - PetscCall(DMGetApplicationContext(dm0,&(ctx0))); - PetscCall(DMGetLocalVector(dm0,&(xloc))); - PetscCall(DMGlobalToLocalBegin(dm0,X,INSERT_VALUES,xloc)); - PetscCall(DMGlobalToLocalEnd(dm0,X,INSERT_VALUES,xloc)); - PetscCall(DMGetLocalVector(dm0,&(yloc))); - PetscCall(VecSet(yloc,0.0)); - PetscCall(VecGetArray(yloc,&y_v_vec)); - PetscCall(VecGetArray(xloc,&x_v_vec)); - PetscCall(DMDAGetLocalInfo(dm0,&(info))); - - PetscScalar (* x_v)[info.gxm] = (PetscScalar (*)[info.gxm]) x_v_vec; - PetscScalar (* y_v)[info.gxm] = (PetscScalar (*)[info.gxm]) y_v_vec; - - PetscInt i, j; - PetscReal xymin[2], xymax[2], hx, hy, darea, scx, scy, scdiag, x, y, - ve, vw, vn, vs, ue, uw, un, us; - - hx = 1./ (info.mx - 1); - hy = 1./ (info.my - 1); - darea = hx * hy; // multiply FD equations by this - scx = 1. / (hx*hx); - scy = 1. / (hy*hy); - scdiag = 2.0 * (scx + scy); // diagonal scaling - for (j = info.ys; j < info.ys + info.ym; j++) { - y = xymin[1] + j * hy; - for (i = info.xs; i < info.xs + info.xm; i++) { - x = xymin[0] + i * hx; - if (i==0 || i==info.mx-1 || j==0 || j==info.my-1) { - y_v[j][i] = x_v[j][i]; - } else { - ve = x_v[j][i+1]; - vw = x_v[j][i-1]; - vn = x_v[j+1][i]; - vs = x_v[j-1][i]; - y_v[j][i] = scdiag * x_v[j][i] - scx * (vw + ve) - scy * (vs + vn); - - } - } - } - - PetscCall(VecRestoreArray(yloc,&y_v_vec)); - PetscCall(VecRestoreArray(xloc,&x_v_vec)); - PetscCall(DMLocalToGlobalBegin(dm0,yloc,ADD_VALUES,Y)); - PetscCall(DMLocalToGlobalEnd(dm0,yloc,ADD_VALUES,Y)); - PetscCall(DMRestoreLocalVector(dm0,&(xloc))); - PetscCall(DMRestoreLocalVector(dm0,&(yloc))); - - PetscFunctionReturn(0); -} - -PetscErrorCode J10_MatMult(Mat J, Vec X, Vec Y) -{ - PetscFunctionBeginUser; - - DM dm0; - DMDALocalInfo info; - Vec xloc; - Vec yloc; - - BiharmCtx * ctx0; - PetscScalar * x_v_vec; - PetscScalar * y_v_vec; - - PetscCall(MatGetDM(J,&(dm0))); - PetscCall(DMGetApplicationContext(dm0,&(ctx0))); - PetscCall(DMGetLocalVector(dm0,&(xloc))); - PetscCall(DMGlobalToLocalBegin(dm0,X,INSERT_VALUES,xloc)); - PetscCall(DMGlobalToLocalEnd(dm0,X,INSERT_VALUES,xloc)); - PetscCall(DMGetLocalVector(dm0,&(yloc))); - PetscCall(VecSet(yloc,0.0)); - PetscCall(VecGetArray(yloc,&y_v_vec)); - PetscCall(VecGetArray(xloc,&x_v_vec)); - PetscCall(DMDAGetLocalInfo(dm0,&(info))); - - PetscScalar (* x_v)[info.gxm] = (PetscScalar (*)[info.gxm]) x_v_vec; - PetscScalar (* y_v)[info.gxm] = (PetscScalar (*)[info.gxm]) y_v_vec; - - PetscInt i, j; - PetscReal xymin[2], xymax[2], hx, hy, darea, scx, scy, scdiag, x, y, - ve, vw, vn, vs, ue, uw, un, us; - - hx = 1. / (info.mx - 1); - hy = 1. / (info.my - 1); - - darea = hx * hy; // multiply FD equations by this - scx = 1. / (hx*hx); - scy = 1. / (hy*hy); - scdiag = 2.0 * (scx + scy); // diagonal scaling - // print info.ys to screen - // PetscCall(PetscPrintf(PETSC_COMM_WORLD,"info.ys = %d\n",info.xm)); - for (j = info.ys+1; j < info.ys + info.ym-1; j++) { - for (i = info.xs+1; i < info.xs + info.xm-1; i++) { - y_v[j][i] = -x_v[j][i]; - } - } - - PetscCall(VecRestoreArray(yloc,&y_v_vec)); - PetscCall(VecRestoreArray(xloc,&x_v_vec)); - PetscCall(DMLocalToGlobalBegin(dm0,yloc,ADD_VALUES,Y)); - PetscCall(DMLocalToGlobalEnd(dm0,yloc,ADD_VALUES,Y)); - PetscCall(DMRestoreLocalVector(dm0,&(xloc))); - PetscCall(DMRestoreLocalVector(dm0,&(yloc))); - - PetscFunctionReturn(0); -} - - -PetscErrorCode J11_MatMult(Mat J, Vec X, Vec Y) -{ - PetscFunctionBeginUser; - - DM dm0; - DMDALocalInfo info; - Vec xloc; - Vec yloc; - - BiharmCtx * ctx0; - PetscScalar * x_v_vec; - PetscScalar * y_v_vec; - - PetscCall(MatGetDM(J,&(dm0))); - PetscCall(DMGetApplicationContext(dm0,&(ctx0))); - PetscCall(DMGetLocalVector(dm0,&(xloc))); - PetscCall(DMGlobalToLocalBegin(dm0,X,INSERT_VALUES,xloc)); - PetscCall(DMGlobalToLocalEnd(dm0,X,INSERT_VALUES,xloc)); - PetscCall(DMGetLocalVector(dm0,&(yloc))); - PetscCall(VecSet(yloc,0.0)); - PetscCall(VecGetArray(yloc,&y_v_vec)); - PetscCall(VecGetArray(xloc,&x_v_vec)); - PetscCall(DMDAGetLocalInfo(dm0,&(info))); - - PetscScalar (* x_v)[info.gxm] = (PetscScalar (*)[info.gxm]) x_v_vec; - PetscScalar (* y_v)[info.gxm] = (PetscScalar (*)[info.gxm]) y_v_vec; - - PetscInt i, j; - PetscReal xymin[2], xymax[2], hx, hy, darea, scx, scy, scdiag, x, y, - ve, vw, vn, vs, ue, uw, un, us; -// PetscCall(DMGetBoundingBox(info.da,xymin,xymax)); - hx = 1. / (info.mx - 1); - hy = 1. / (info.my - 1); - - darea = hx * hy; - scx = 1. / (hx*hx); - scy = 1. / (hy*hy); - scdiag = 2.0 * (scx + scy); - for (j = info.ys; j < info.ys + info.ym; j++) { - y = xymin[1] + j * hy; - for (i = info.xs; i < info.xs + info.xm; i++) { - x = xymin[0] + i * hx; - if (i==0 || i==info.mx-1 || j==0 || j==info.my-1) { - y_v[j][i] = x_v[j][i]; - } else { - ve = x_v[j][i+1]; - vw = x_v[j][i-1]; - vn = x_v[j+1][i]; - vs = x_v[j-1][i]; - y_v[j][i] = scdiag * x_v[j][i] - scx * (vw + ve) - scy * (vs + vn); - - } - } - } - - PetscCall(VecRestoreArray(yloc,&y_v_vec)); - PetscCall(VecRestoreArray(xloc,&x_v_vec)); - PetscCall(DMLocalToGlobalBegin(dm0,yloc,ADD_VALUES,Y)); - PetscCall(DMLocalToGlobalEnd(dm0,yloc,ADD_VALUES,Y)); - PetscCall(DMRestoreLocalVector(dm0,&(xloc))); - PetscCall(DMRestoreLocalVector(dm0,&(yloc))); - - PetscFunctionReturn(0); -} - -PetscErrorCode WholeMatMult(Mat J, Vec X, Vec Y) -{ - Vec J00X; - Vec J00Y; - Vec J10X; - Vec J10Y; - Vec J11X; - Vec J11Y; - - struct SubMatrixCtx * J00ctx; - struct SubMatrixCtx * J10ctx; - struct SubMatrixCtx * J11ctx; - struct JacobianCtx * jctx; - - PetscFunctionBeginUser; - - PetscCall(MatShellGetContext(J,&(jctx))); - - PetscCall(VecSet(Y,0.0)); - - Mat J00 = jctx->submats[0]; - PetscCall(MatShellGetContext(J00,&(J00ctx))); - PetscCall(VecGetSubVector(X,*(J00ctx->cols),&(J00X))); - PetscCall(VecGetSubVector(Y,*(J00ctx->rows),&(J00Y))); - PetscCall(MatMult(J00,J00X,J00Y)); - PetscCall(VecRestoreSubVector(X,*(J00ctx->cols),&(J00X))); - PetscCall(VecRestoreSubVector(Y,*(J00ctx->rows),&(J00Y))); - - Mat J10 = jctx->submats[2]; - PetscCall(MatShellGetContext(J10,&(J10ctx))); - PetscCall(VecGetSubVector(X,*(J10ctx->cols),&(J10X))); - PetscCall(VecGetSubVector(Y,*(J10ctx->rows),&(J10Y))); - PetscCall(MatMult(J10,J10X,J10Y)); - PetscCall(VecRestoreSubVector(X,*(J10ctx->cols),&(J10X))); - PetscCall(VecRestoreSubVector(Y,*(J10ctx->rows),&(J10Y))); - - Mat J11 = jctx->submats[3]; - PetscCall(MatShellGetContext(J11,&(J11ctx))); - PetscCall(VecGetSubVector(X,*(J11ctx->cols),&(J11X))); - PetscCall(VecGetSubVector(Y,*(J11ctx->rows),&(J11Y))); - PetscCall(MatMult(J11,J11X,J11Y)); - PetscCall(VecRestoreSubVector(X,*(J11ctx->cols),&(J11X))); - PetscCall(VecRestoreSubVector(Y,*(J11ctx->rows),&(J11Y))); - - - PetscFunctionReturn(0); -} - - -PetscErrorCode MatCreateSubMatrices0(Mat J, PetscInt nfields, IS * irow, IS * icol, MatReuse scall, Mat * * submats) -{ - PetscFunctionBeginUser; - - PetscInt M; - PetscInt N; - Mat block; - DM dm0; - PetscInt dof; - - struct UserCtx0 * ctx0; - struct JacobianCtx * jctx; - struct SubMatrixCtx * subctx; - - PetscCall(MatShellGetContext(J,&(jctx))); - DM * subdms = jctx->subdms; - - PetscInt nsubmats = nfields*nfields; - PetscCall(PetscCalloc1(nsubmats,submats)); - PetscCall(MatGetDM(J,&(dm0))); - PetscCall(DMGetApplicationContext(dm0,&(ctx0))); - PetscCall(DMDAGetInfo(dm0,NULL,&(M),&(N),NULL,NULL,NULL,NULL,&(dof),NULL,NULL,NULL,NULL,NULL)); - PetscInt subblockrows = M*N; - PetscInt subblockcols = M*N; - Mat * submat_arr = *submats; - - for (int i = 0; i <= nsubmats - 1; i += 1) - { - PetscCall(MatCreate(PETSC_COMM_WORLD,&(block))); - PetscCall(MatSetSizes(block,PETSC_DECIDE,PETSC_DECIDE,subblockrows,subblockcols)); - PetscCall(MatSetType(block,MATSHELL)); - PetscCall(PetscMalloc1(1,&(subctx))); - PetscInt rowidx = i / dof; - PetscInt colidx = (i)%(dof); - subctx->rows = &(irow[rowidx]); - subctx->cols = &(icol[colidx]); - PetscCall(DMSetApplicationContext(subdms[rowidx],ctx0)); - PetscCall(MatSetDM(block,subdms[rowidx])); - PetscCall(MatShellSetContext(block,subctx)); - PetscCall(MatSetUp(block)); - submat_arr[i] = block; - } - PetscCall(MatShellSetOperation(submat_arr[0],MATOP_MULT,(void (*)(void))J00_MatMult)); - PetscCall(MatShellSetOperation(submat_arr[2],MATOP_MULT,(void (*)(void))J10_MatMult)); - PetscCall(MatShellSetOperation(submat_arr[3],MATOP_MULT,(void (*)(void))J11_MatMult)); - - PetscFunctionReturn(0); -} - - -PetscErrorCode PopulateMatContext(struct JacobianCtx * jctx, DM * subdms, IS * fields) -{ - PetscFunctionBeginUser; - - jctx->subdms = subdms; - jctx->fields = fields; - - PetscFunctionReturn(0); -} diff --git a/tests/test_petsc.py b/tests/test_petsc.py index 8f9cc4d9a8..e4cc05e024 100644 --- a/tests/test_petsc.py +++ b/tests/test_petsc.py @@ -1,6 +1,5 @@ import numpy as np import os -import pytest from conftest import skipif from devito import (Grid, Function, TimeFunction, Eq, Operator, From 8386992147685e923964601c87e11560e26bce1c Mon Sep 17 00:00:00 2001 From: ZoeLeibowitz Date: Thu, 5 Jun 2025 16:35:49 +0100 Subject: [PATCH 39/84] dsl: Improve Jacobian abstraction --- devito/petsc/iet/routines.py | 57 +++---- devito/petsc/solve.py | 6 +- devito/petsc/types/equation.py | 4 +- devito/petsc/types/types.py | 223 ++++++++++++------------- examples/petsc/random/02_biharmonic.py | 3 +- tests/test_petsc.py | 39 ++--- 6 files changed, 164 insertions(+), 168 deletions(-) diff --git a/devito/petsc/iet/routines.py b/devito/petsc/iet/routines.py index 8f580a1b52..8907cfe795 100644 --- a/devito/petsc/iet/routines.py +++ b/devito/petsc/iet/routines.py @@ -111,7 +111,7 @@ def _make_core(self): self._make_matvec(fielddata.arrays, fielddata.jacobian.matvecs) self._make_formfunc(fielddata) self._make_formrhs(fielddata) - if fielddata.initialguess.equations: + if fielddata.initialguess.eqs: self._make_initialguess(fielddata) self._make_user_struct_callback() @@ -515,7 +515,7 @@ def _create_form_rhs_body(self, body, fielddata): return Uxreplace(subs).visit(formrhs_body) def _make_initialguess(self, fielddata): - initguess = fielddata.initialguess.equations + initguess = fielddata.initialguess.eqs sobjs = self.solver_objs # Compile initital guess `eqns` into an IET via recursive compilation @@ -678,15 +678,11 @@ def zero_memory(self): def _make_core(self): injectsolve = self.injectsolve - targets = injectsolve.expr.rhs.fielddata.targets all_fielddata = injectsolve.expr.rhs.fielddata - for t in targets: - row_matvecs = all_fielddata.jacobian.submatrices[t] - arrays = all_fielddata.arrays[t] - for submat, mtvs in row_matvecs.items(): - if mtvs['matvecs']: - self._make_matvec(arrays, mtvs['matvecs'], prefix=f'{submat}_MatMult') + for sm in all_fielddata.jacobian.nonzero_submatrices: + arrays = all_fielddata.arrays[sm.row_target] + self._make_matvec(arrays, sm.matvecs, prefix=f'{sm.name}_MatMult') self._make_whole_matvec() self._make_whole_formfunc(all_fielddata) @@ -714,7 +710,7 @@ def _whole_matvec_body(self): jctx = objs['ljacctx'] ctx_main = petsc_call('MatShellGetContext', [objs['J'], Byref(jctx)]) - nonzero_submats = self.jacobian.nonzero_submatrix_keys + nonzero_submats = self.jacobian.nonzero_submatrices zero_y_memory = petsc_call( 'VecSet', [objs['Y'], 0.0] @@ -722,17 +718,17 @@ def _whole_matvec_body(self): calls = () for sm in nonzero_submats: - idx = self.jacobian.submat_to_index[sm] - ctx = sobjs[f'{sm}ctx'] - X = sobjs[f'{sm}X'] - Y = sobjs[f'{sm}Y'] + name = sm.name + ctx = sobjs[f'{name}ctx'] + X = sobjs[f'{name}X'] + Y = sobjs[f'{name}Y'] rows = objs['rows'].base cols = objs['cols'].base - sm_indexed = objs['Submats'].indexed[idx] + sm_indexed = objs['Submats'].indexed[sm.linear_idx] calls += ( - DummyExpr(sobjs[sm], FieldFromPointer(sm_indexed, jctx)), - petsc_call('MatShellGetContext', [sobjs[sm], Byref(ctx)]), + DummyExpr(sobjs[name], FieldFromPointer(sm_indexed, jctx)), + petsc_call('MatShellGetContext', [sobjs[name], Byref(ctx)]), petsc_call( 'VecGetSubVector', [objs['X'], Deref(FieldFromPointer(cols, ctx)), Byref(X)] @@ -741,7 +737,7 @@ def _whole_matvec_body(self): 'VecGetSubVector', [objs['Y'], Deref(FieldFromPointer(rows, ctx)), Byref(Y)] ), - petsc_call('MatMult', [sobjs[sm], X, Y]), + petsc_call('MatMult', [sobjs[name], X, Y]), petsc_call( 'VecRestoreSubVector', [objs['X'], Deref(FieldFromPointer(cols, ctx)), Byref(X)] @@ -1012,19 +1008,19 @@ def _submat_callback_body(self): upper_bound = objs['nsubmats'] - 1 iteration = Iteration(List(body=iter_body), i, upper_bound) - nonzero_submats = self.jacobian.nonzero_submatrix_keys + nonzero_submats = self.jacobian.nonzero_submatrices matvec_lookup = {mv.name.split('_')[0]: mv for mv in self.matvecs} matmult_op = [ petsc_call( 'MatShellSetOperation', [ - objs['submat_arr'].indexed[self.jacobian.submat_to_index[sb]], + objs['submat_arr'].indexed[sb.linear_idx], 'MATOP_MULT', - MatShellSetOp(matvec_lookup[sb].name, void, void), + MatShellSetOp(matvec_lookup[sb.name].name, void, void), ], ) - for sb in nonzero_submats if sb in matvec_lookup + for sb in nonzero_submats if sb.name in matvec_lookup ] body = [ @@ -1158,22 +1154,23 @@ def _extend_build(self, base_dict): }) jacobian = injectsolve.expr.rhs.fielddata.jacobian - submatrix_keys = jacobian.submatrix_keys + submatrices = jacobian.nonzero_submatrices base_dict['jacctx'] = JacobianStruct( name=sreg.make_name(prefix=objs['ljacctx'].name), fields=objs['ljacctx'].fields, ) - for key in submatrix_keys: - base_dict[key] = Mat(name=key) - base_dict[f'{key}ctx'] = SubMatrixStruct( - name=f'{key}ctx', + for sm in submatrices: + name = sm.name + base_dict[name] = Mat(name=name) + base_dict[f'{name}ctx'] = SubMatrixStruct( + name=f'{name}ctx', fields=objs['subctx'].fields, ) - base_dict[f'{key}X'] = CallbackVec(f'{key}X') - base_dict[f'{key}Y'] = CallbackVec(f'{key}Y') - base_dict[f'{key}F'] = CallbackVec(f'{key}F') + base_dict[f'{name}X'] = CallbackVec(f'{name}X') + base_dict[f'{name}Y'] = CallbackVec(f'{name}Y') + base_dict[f'{name}F'] = CallbackVec(f'{name}F') # Bundle objects/metadata required by the coupled residual callback f_components = [] diff --git a/devito/petsc/solve.py b/devito/petsc/solve.py index d2f0f39809..7e4f02df12 100644 --- a/devito/petsc/solve.py +++ b/devito/petsc/solve.py @@ -86,9 +86,7 @@ def linear_solve_args(self): eqns = sorted(eqns, key=lambda e: 0 if isinstance(e, EssentialBC) else 1) jacobian = Jacobian(target, eqns, arrays, self.time_mapper) - residual = Residual(target, eqns, arrays, self.time_mapper, jacobian.scdiag) - initialguess = InitialGuess(target, eqns, arrays) field_data = FieldData( @@ -129,8 +127,8 @@ def linear_solve_args(self): ) residual = MixedResidual( - self.target_eqns, arrays, - self.time_mapper, jacobian.target_scaler_mapper + self.target_eqns, arrays, self.time_mapper, + jacobian.target_scaler_mapper ) all_data = MultipleFieldData( diff --git a/devito/petsc/types/equation.py b/devito/petsc/types/equation.py index ad929c8ed5..6d476c12f7 100644 --- a/devito/petsc/types/equation.py +++ b/devito/petsc/types/equation.py @@ -26,7 +26,7 @@ class ZeroRow(EssentialBC): """ Equation used to zero all entries, except the diagonal, of a row in the Jacobian. - + Note: This is only used interally by the compiler, not by users. """ @@ -36,7 +36,7 @@ class ZeroRow(EssentialBC): class ZeroColumn(EssentialBC): """ Equation used to zero the column of the Jacobian. - + Note: This is only used interally by the compiler, not by users. """ diff --git a/devito/petsc/types/types.py b/devito/petsc/types/types.py index a0a9cd2fe5..8bc08d804d 100644 --- a/devito/petsc/types/types.py +++ b/devito/petsc/types/types.py @@ -235,6 +235,16 @@ def targets(self): class Jacobian: + """ + Represents a Jacobian matrix in a matrix-free form. + + This Jacobian is defined implicitly via matrix-vector products + derived from the symbolic equations provided in `matvecs`. + + It assumes the problem is linear, meaning the Jacobian + corresponds to a constant coefficient matrix and does not + require explicit symbolic differentiation. + """ def __init__(self, target, eqns, arrays, time_mapper): self.target = target self.eqns = eqns @@ -318,140 +328,129 @@ def _scale_bcs(self, matvecs, scdiag): ] +class SubMatrixBlock: + def __init__(self, name, matvecs, scdiag, + row_target, col_target, row_idx, col_idx, linear_idx): + self.name = name + self.matvecs = matvecs + self.scdiag = scdiag + self.row_target = row_target + self.col_target = col_target + self.row_idx = row_idx + self.col_idx = col_idx + self.linear_idx = linear_idx + + def is_diag(self): + return self.row_idx == self.col_idx + + def __repr__(self): + return (f"") + + class MixedJacobian(Jacobian): + """ + Represents a Jacobian for a linear system with a solution vector + composed of multiple fields (targets). + + Defines matrix-vector products for each sub-block, + each with its own generated callback. The matrix may be treated + as monolithic or block-structured in PETSc, but sub-block + callbacks are generated in both cases. + + Assumes a **linear** problem, so this Jacobian corresponds to a + coefficient matrix and does not require differentiation. + + # TODO: pcfieldsplit support for each block + """ def __init__(self, target_eqns, arrays, time_mapper): """ """ self.targets = as_tuple(target_eqns.keys()) self.arrays = arrays self.time_mapper = time_mapper - self.submatrices = self._initialize_submatrices() + self._submatrices = [] self._build_blocks(target_eqns) - def _initialize_submatrices(self): - """ - Create a dict of submatrices for each target with metadata. - """ - submatrices = {} - num_targets = len(self.targets) - - for i, target in enumerate(self.targets): - submatrices[target] = {} - for j in range(num_targets): - key = f'J{i}{j}' - submatrices[target][key] = { - 'matvecs': None, - 'derivative_wrt': self.targets[j], - 'index': i * num_targets + j, - 'scdiag': None - } - - return submatrices - - def _build_blocks(self, target_eqns): - for target, eqns in target_eqns.items(): - self._build_block(target, eqns) - - def _build_block(self, target, eqns): - arrays = self.arrays[target] - for submat, mtvs in self.submatrices[target].items(): - matvecs = [ - e for eq in eqns for e in - self._build_matvec_eq(eq, mtvs['derivative_wrt'], arrays) - ] - matvecs = [m for m in matvecs if m is not None] - matvecs = tuple(sorted(matvecs, key=lambda e: not isinstance(e, EssentialBC))) - - matvecs = self._scale_non_bcs(matvecs, target) - scdiag = self._compute_scdiag(matvecs, arrays) - matvecs = self._scale_bcs(matvecs, scdiag) - - if matvecs: - self.set_submatrix(target, submat, matvecs, scdiag) - @property - def submatrix_keys(self): + def submatrices(self): """ - Return a list of all submatrix keys (e.g., ['J00', 'J01', 'J10', 'J11']). + Return a list of all submatrix blocks. + Each block contains metadata about the matrix-vector products. """ - return [key for submats in self.submatrices.values() for key in submats.keys()] + return self._submatrices @property - def nonzero_submatrix_keys(self): + def no_submatrices(self): """ - Returns a list of submats where 'matvecs' is not None. + Return the number of submatrix blocks. """ - return [ - key - for submats in self.submatrices.values() - for key, value in submats.items() - if value['matvecs'] is not None - ] + return len(self._submatrices) @property - def submat_to_index(self): - """ - Returns a dict mapping submatrix keys to their index. - """ - return { - key: value['index'] - for submats in self.submatrices.values() - for key, value in submats.items() - } - - # CHECK/TEST THIS - def is_diagonal_submatrix(self, key): - """ - Return True if the given key corresponds to a diagonal - submatrix (e.g., 'J00', 'J11'), else False. - """ - for i, t in enumerate(self.targets): - diag_key = f'J{i}{i}' - if key == diag_key and diag_key in self.submatrices[t]: - return True - return False - - def set_submatrix(self, field, key, matvecs, scdiag): - """ - Set the matrix-vector equations for a submatrix. - - Parameters - ---------- - field : Function - The target field that the submatrix operates on. - key: str - The identifier for the submatrix (e.g., 'J00', 'J01'). - matvecs: list of Eq - The matrix-vector equations forming the submatrix. - """ - if field in self.submatrices and key in self.submatrices[field]: - self.submatrices[field][key]['matvecs'] = matvecs - self.submatrices[field][key]['scdiag'] = scdiag - else: - raise KeyError(f'Invalid field ({field}) or submatrix key ({key})') - - def get_submatrix(self, field, key): - """ - Retrieve a specific submatrix. - """ - return self.submatrices.get(field, {}).get(key, None) + def nonzero_submatrices(self): + """Return SubMatrixBlock objects that have non-empty matvecs.""" + return [submat for submat in self.submatrices if submat.matvecs] @property def target_scaler_mapper(self): """ - Return a mapping from each target to its diagonal submatrix's scaler. + Map each row_target index to its diagonal submatrix's scaler. """ mapper = {} - for target, submats in self.submatrices.items(): - for key, value in submats.items(): - if self.is_diagonal_submatrix(key): - mapper[target] = value.get('scdiag') - break + for sm in self.submatrices: + if sm.row_idx == sm.col_idx: + mapper[sm.row_target] = sm.scdiag return mapper + def _build_blocks(self, target_eqns): + """ + Build all SubMatrixBlock objects for the Jacobian. + """ + for i, row_target in enumerate(self.targets): + eqns = target_eqns[row_target] + arrays = self.arrays[row_target] + for j, col_target in enumerate(self.targets): + matvecs = [ + e for eq in eqns for e in + self._build_matvec_eq(eq, col_target, arrays) + ] + matvecs = [m for m in matvecs if m is not None] + # Sort to put EssentialBC first if any + matvecs = tuple( + sorted(matvecs, key=lambda e: not isinstance(e, EssentialBC)) + ) + matvecs = self._scale_non_bcs(matvecs, row_target) + scdiag = self._compute_scdiag(matvecs, arrays) + matvecs = self._scale_bcs(matvecs, scdiag) + + name = f'J{i}{j}' + block = SubMatrixBlock( + name=name, + matvecs=matvecs, + scdiag=scdiag, + row_target=row_target, + col_target=col_target, + row_idx=i, + col_idx=j, + linear_idx=i * len(self.targets) + j + ) + self._submatrices.append(block) + + def get_submatrix(self, row_idx, col_idx): + """ + Return the SubMatrixBlock at (row_idx, col_idx), or None if not found. + """ + for sm in self.submatrices: + if sm.row_idx == row_idx and sm.col_idx == col_idx: + return sm + return None + def __repr__(self): - # TODO: edit - return str(self.submatrices) + summary = ', '.join( + f"{sm.name} (row={sm.row_idx}, col={sm.col_idx})" + for sm in self.submatrices + ) + return f"" class Residual: @@ -486,7 +485,7 @@ def _build_equations(self): for eq in self.eqns: b, F_target, _, targets = separate_eqn(eq, self.target) funcs.extend(self._make_F_target(eq, F_target, targets)) - # TODO: if b is zero then don't need a rhs vector+callback + # TODO: If b is zero then don't need a rhs vector+callback rhs.extend(self._make_b(eq, b)) self._formfuncs = [self._scale_bcs(eq) for eq in funcs] @@ -596,18 +595,18 @@ def __init__(self, target, eqns, arrays): self.eqns = as_tuple(eqns) self.arrays = arrays self._build_equations() - + @property - def equations(self): + def eqs(self): """ """ - return self._equations + return self._eqs def _build_equations(self): """ Return a list of initial guess equations. """ - self._equations = [ + self._eqs = [ eq for eq in (self._make_initial_guess(e) for e in self.eqns) if eq is not None diff --git a/examples/petsc/random/02_biharmonic.py b/examples/petsc/random/02_biharmonic.py index 973b16499e..214e38afda 100644 --- a/examples/petsc/random/02_biharmonic.py +++ b/examples/petsc/random/02_biharmonic.py @@ -129,7 +129,8 @@ def f_fcn(x, y): with switchconfig(language='petsc'): op = Operator(petsc) - op.apply() + # op.apply() + print(op.ccode) u_exact.data[:] = u_exact_fcn(X, Y) lap_u.data[:] = lap_u_exact_fcn(X, Y) diff --git a/tests/test_petsc.py b/tests/test_petsc.py index e4cc05e024..d0f4bdd087 100644 --- a/tests/test_petsc.py +++ b/tests/test_petsc.py @@ -817,36 +817,37 @@ def test_submatrices(self): jacobian = petsc[0].rhs.fielddata.jacobian - j00 = jacobian.get_submatrix(e, 'J00') - j01 = jacobian.get_submatrix(e, 'J01') - j10 = jacobian.get_submatrix(g, 'J10') - j11 = jacobian.get_submatrix(g, 'J11') + j00 = jacobian.get_submatrix(0, 0) + j01 = jacobian.get_submatrix(0, 1) + j10 = jacobian.get_submatrix(1, 0) + j11 = jacobian.get_submatrix(1, 1) # Check the number of submatrices - assert len(jacobian.submatrix_keys) == 4 - assert str(jacobian.submatrix_keys) == "['J00', 'J01', 'J10', 'J11']" + assert jacobian.no_submatrices == 4 # Technically a non-coupled problem, so the only non-zero submatrices # should be the diagonal ones i.e J00 and J11 - assert jacobian.nonzero_submatrix_keys == ['J00', 'J11'] - assert jacobian.get_submatrix(e, 'J01')['matvecs'] is None - assert jacobian.get_submatrix(g, 'J10')['matvecs'] is None - - j00 = jacobian.get_submatrix(e, 'J00') - j11 = jacobian.get_submatrix(g, 'J11') + nonzero_submats = jacobian.nonzero_submatrices + assert len(nonzero_submats) == 2 + assert j00 in nonzero_submats + assert j11 in nonzero_submats + assert j01 not in nonzero_submats + assert j10 not in nonzero_submats + assert not j01.matvecs + assert not j10.matvecs # Compatible scaling to reduce condition number of jacobian - assert str(j00['matvecs'][0]) == 'Eq(y_e(x, y),' \ + assert str(j00.matvecs[0]) == 'Eq(y_e(x, y),' \ + ' h_x*h_y*(Derivative(x_e(x, y), (x, 2)) + Derivative(x_e(x, y), (y, 2))))' - assert str(j11['matvecs'][0]) == 'Eq(y_g(x, y),' \ + assert str(j11.matvecs[0]) == 'Eq(y_g(x, y),' \ + ' h_x*h_y*(Derivative(x_g(x, y), (x, 2)) + Derivative(x_g(x, y), (y, 2))))' - # Check the derivative wrt fields - assert j00['derivative_wrt'] == e - assert j01['derivative_wrt'] == g - assert j10['derivative_wrt'] == e - assert j11['derivative_wrt'] == g + # Check the col_targets + assert j00.col_target == e + assert j01.col_target == g + assert j10.col_target == e + assert j11.col_target == g @skipif('petsc') def test_residual_bundle(self): From fb207810eb23a24c4784fa20f97c89d77a178ba3 Mon Sep 17 00:00:00 2001 From: ZoeLeibowitz Date: Thu, 5 Jun 2025 18:01:40 +0100 Subject: [PATCH 40/84] examples: Add pressure norm check in stokes example --- devito/petsc/iet/routines.py | 81 ++++---- devito/petsc/solve.py | 9 +- devito/petsc/types/array.py | 17 +- devito/petsc/types/types.py | 57 +++--- devito/types/array.py | 6 +- examples/petsc/Poisson/01_poisson.py | 1 + examples/petsc/cfd/01_navierstokes.py | 34 ++-- tests/test_petsc.py | 256 +++++++++++++++++++++----- 8 files changed, 333 insertions(+), 128 deletions(-) diff --git a/devito/petsc/iet/routines.py b/devito/petsc/iet/routines.py index 8907cfe795..2fbdafe1ef 100644 --- a/devito/petsc/iet/routines.py +++ b/devito/petsc/iet/routines.py @@ -106,22 +106,30 @@ def zero_memory(self): in the callback.""" return True + @property + def fielddata(self): + return self.injectsolve.expr.rhs.fielddata + + @property + def arrays(self): + return self.fielddata.arrays + def _make_core(self): - fielddata = self.injectsolve.expr.rhs.fielddata - self._make_matvec(fielddata.arrays, fielddata.jacobian.matvecs) - self._make_formfunc(fielddata) - self._make_formrhs(fielddata) - if fielddata.initialguess.eqs: - self._make_initialguess(fielddata) + self._make_matvec(self.fielddata.jacobian) + self._make_formfunc() + self._make_formrhs() + if self.fielddata.initialguess.eqs: + self._make_initialguess() self._make_user_struct_callback() - def _make_matvec(self, arrays, matvecs, prefix='MatMult'): + def _make_matvec(self, jacobian, prefix='MatMult'): # Compile matvec `eqns` into an IET via recursive compilation + matvecs = jacobian.matvecs irs_matvec, _ = self.rcompile(matvecs, options={'mpi': False}, sregistry=self.sregistry, concretize_mapper=self.concretize_mapper) body_matvec = self._create_matvec_body(List(body=irs_matvec.uiet.body), - arrays) + jacobian) objs = self.objs cb = PETScCallable( @@ -133,7 +141,7 @@ def _make_matvec(self, arrays, matvecs, prefix='MatMult'): self._matvecs.append(cb) self._efuncs[cb.name] = cb - def _create_matvec_body(self, body, arrays): + def _create_matvec_body(self, body, jacobian): linsolve_expr = self.injectsolve.expr.rhs objs = self.objs sobjs = self.solver_objs @@ -142,8 +150,8 @@ def _create_matvec_body(self, body, arrays): ctx = objs['dummyctx'] xlocal = objs['xloc'] ylocal = objs['yloc'] - y_matvec = arrays['y'] - x_matvec = arrays['x'] + y_matvec = self.arrays[jacobian.row_target]['y'] + x_matvec = self.arrays[jacobian.col_target]['x'] body = self.timedep.uxreplace_time(body) @@ -265,15 +273,15 @@ def _create_matvec_body(self, body, arrays): self._struct_params.extend(fields) return matvec_body - def _make_formfunc(self, fielddata): - formfuncs = fielddata.residual.formfuncs + def _make_formfunc(self): + formfuncs = self.fielddata.residual.formfuncs # Compile formfunc `eqns` into an IET via recursive compilation irs_formfunc, _ = self.rcompile( formfuncs, options={'mpi': False}, sregistry=self.sregistry, concretize_mapper=self.concretize_mapper ) body_formfunc = self._create_formfunc_body( - List(body=irs_formfunc.uiet.body), fielddata + List(body=irs_formfunc.uiet.body) ) objs = self.objs cb = PETScCallable( @@ -285,10 +293,11 @@ def _make_formfunc(self, fielddata): self._formfuncs.append(cb) self._efuncs[cb.name] = cb - def _create_formfunc_body(self, body, fielddata): + def _create_formfunc_body(self, body): linsolve_expr = self.injectsolve.expr.rhs objs = self.objs sobjs = self.solver_objs + target = self.fielddata.target dmda = sobjs['callbackdm'] ctx = objs['dummyctx'] @@ -298,8 +307,8 @@ def _create_formfunc_body(self, body, fielddata): fields = self._dummy_fields(body) self._struct_params.extend(fields) - f_formfunc = fielddata.arrays['f'] - x_formfunc = fielddata.arrays['x'] + f_formfunc = self.fielddata.arrays[target]['f'] + x_formfunc = self.fielddata.arrays[target]['x'] dm_cast = DummyExpr(dmda, DMCast(objs['dummyptr']), init=True) @@ -402,8 +411,8 @@ def _create_formfunc_body(self, body, fielddata): return Uxreplace(subs).visit(formfunc_body) - def _make_formrhs(self, fielddata): - formrhs = fielddata.residual.formrhs + def _make_formrhs(self): + formrhs = self.fielddata.residual.formrhs sobjs = self.solver_objs # Compile formrhs `eqns` into an IET via recursive compilation @@ -412,7 +421,7 @@ def _make_formrhs(self, fielddata): concretize_mapper=self.concretize_mapper ) body_formrhs = self._create_form_rhs_body( - List(body=irs_formrhs.uiet.body), fielddata + List(body=irs_formrhs.uiet.body) ) objs = self.objs cb = PETScCallable( @@ -424,10 +433,11 @@ def _make_formrhs(self, fielddata): self._formrhs.append(cb) self._efuncs[cb.name] = cb - def _create_form_rhs_body(self, body, fielddata): + def _create_form_rhs_body(self, body): linsolve_expr = self.injectsolve.expr.rhs objs = self.objs sobjs = self.solver_objs + target = self.fielddata.target dmda = sobjs['callbackdm'] ctx = objs['dummyctx'] @@ -446,7 +456,7 @@ def _create_form_rhs_body(self, body, fielddata): sobjs['blocal'] ]) - b_arr = fielddata.arrays['b'] + b_arr = self.fielddata.arrays[target]['b'] vec_get_array = petsc_call( 'VecGetArray', [sobjs['blocal'], Byref(b_arr._C_symbol)] @@ -514,8 +524,8 @@ def _create_form_rhs_body(self, body, fielddata): return Uxreplace(subs).visit(formrhs_body) - def _make_initialguess(self, fielddata): - initguess = fielddata.initialguess.eqs + def _make_initialguess(self): + initguess = self.fielddata.initialguess.eqs sobjs = self.solver_objs # Compile initital guess `eqns` into an IET via recursive compilation @@ -524,7 +534,7 @@ def _make_initialguess(self, fielddata): concretize_mapper=self.concretize_mapper ) body_init_guess = self._create_initial_guess_body( - List(body=irs.uiet.body), fielddata + List(body=irs.uiet.body) ) objs = self.objs cb = PETScCallable( @@ -536,15 +546,16 @@ def _make_initialguess(self, fielddata): self._initialguesses.append(cb) self._efuncs[cb.name] = cb - def _create_initial_guess_body(self, body, fielddata): + def _create_initial_guess_body(self, body): linsolve_expr = self.injectsolve.expr.rhs objs = self.objs sobjs = self.solver_objs + target = self.fielddata.target dmda = sobjs['callbackdm'] ctx = objs['dummyctx'] - x_arr = fielddata.arrays['x'] + x_arr = self.fielddata.arrays[target]['x'] vec_get_array = petsc_call( 'VecGetArray', [objs['xloc'], Byref(x_arr._C_symbol)] @@ -681,8 +692,8 @@ def _make_core(self): all_fielddata = injectsolve.expr.rhs.fielddata for sm in all_fielddata.jacobian.nonzero_submatrices: - arrays = all_fielddata.arrays[sm.row_target] - self._make_matvec(arrays, sm.matvecs, prefix=f'{sm.name}_MatMult') + # arrays = all_fielddata.arrays[sm.row_target] + self._make_matvec(sm, prefix=f'{sm.name}_MatMult') self._make_whole_matvec() self._make_whole_formfunc(all_fielddata) @@ -1044,18 +1055,18 @@ def _submat_callback_body(self): ) def residual_bundle(self, body, bundles): - mapper1 = bundles['bundle_mapper'] + mapper = bundles['bundle_mapper'] indexeds = FindSymbols('indexeds').visit(body) + subs = {} - subss = {} for i in indexeds: - if i.base in mapper1: - bundle = mapper1[i.base] + if i.base in mapper: + bundle = mapper[i.base] index = bundles['target_indices'][i.function.target] index = (index,)+i.indices - subss[i] = bundle.__getitem__(index) + subs[i] = bundle.__getitem__(index) - body = Uxreplace(subss).visit(body) + body = Uxreplace(subs).visit(body) return body diff --git a/devito/petsc/solve.py b/devito/petsc/solve.py index 7e4f02df12..8adb9cda1e 100644 --- a/devito/petsc/solve.py +++ b/devito/petsc/solve.py @@ -81,7 +81,7 @@ def linear_solve_args(self): funcs = get_funcs(eqns) self.time_mapper = generate_time_mapper(funcs) - arrays = self.generate_arrays(target) + arrays = self.generate_arrays_combined(target) eqns = sorted(eqns, key=lambda e: 0 if isinstance(e, EssentialBC) else 1) @@ -108,6 +108,9 @@ def generate_arrays(self, target): for p in prefixes } + def generate_arrays_combined(self, *targets): + return {target: self.generate_arrays(target) for target in targets} + class InjectMixedSolve(InjectSolve): @@ -140,8 +143,8 @@ def linear_solve_args(self): return coupled_targets[0], tuple(funcs), all_data - def generate_arrays_combined(self, *targets): - return {target: self.generate_arrays(target) for target in targets} + # def generate_arrays_combined(self, *targets): + # return {target: self.generate_arrays(target) for target in targets} def generate_time_mapper(funcs): diff --git a/devito/petsc/types/array.py b/devito/petsc/types/array.py index 8b57aca44c..c7a09f8265 100644 --- a/devito/petsc/types/array.py +++ b/devito/petsc/types/array.py @@ -1,3 +1,5 @@ +from sympy import Expr + from functools import cached_property from ctypes import POINTER, Structure @@ -164,7 +166,7 @@ def __getitem__(self, index): elif len(index) == self.ndim + 1: component_index, indices = index[0], index[1:] names = tuple(i.target.name for i in self.components) - return ComponentAccess( + return PetscComponentAccess( self.indexed[indices], component_index, component_names=names @@ -178,6 +180,19 @@ def pname(self): return self._pname +class PetscComponentAccess(ComponentAccess): + def __new__(cls, arg, index=0, component_names=None, **kwargs): + if not arg.is_Indexed: + raise ValueError("Expected Indexed, got `%s` instead" % type(arg)) + names = component_names or cls._default_component_names + + obj = Expr.__new__(cls, arg) + obj._index = index + obj._component_names = names + + return obj + + class AoSIndexedData(IndexedData): @property def dtype(self): diff --git a/devito/petsc/types/types.py b/devito/petsc/types/types.py index 8bc08d804d..21433b5944 100644 --- a/devito/petsc/types/types.py +++ b/devito/petsc/types/types.py @@ -260,6 +260,14 @@ def matvecs(self): def scdiag(self): return self._scdiag + @property + def row_target(self): + return self.target + + @property + def col_target(self): + return self.target + def _build_matvecs(self): matvecs = [ e for eq in self.eqns for e in @@ -275,30 +283,35 @@ def _build_matvecs(self): self._matvecs = matvecs self._scdiag = scdiag - def _build_matvec_eq(self, eq, target=None, arrays=None): - target = target or self.target - arrays = arrays or self.arrays + def _build_matvec_eq(self, eq, col_target=None, row_target=None): + col_target = col_target or self.target + row_target = row_target or self.target - b, F_target, _, targets = separate_eqn(eq, target) + b, F_target, _, targets = separate_eqn(eq, col_target) if F_target: - return self._make_matvec(eq, F_target, targets, arrays) + return self._make_matvec( + eq, F_target, targets, col_target, row_target + ) return (None,) - def _make_matvec(self, eq, F_target, targets, arrays): + def _make_matvec(self, eq, F_target, targets, col_target, row_target): + y = self.arrays[row_target]['y'] + x = self.arrays[col_target]['x'] + if isinstance(eq, EssentialBC): # NOTE: Essential BCs are trivial equations in the solver. # See `EssentialBC` for more details. - rhs = arrays['x'] - zero_row = ZeroRow(arrays['y'], rhs, subdomain=eq.subdomain) - zero_column = ZeroColumn(arrays['x'], 0., subdomain=eq.subdomain) + zero_row = ZeroRow(y, x, subdomain=eq.subdomain) + zero_column = ZeroColumn(x, 0., subdomain=eq.subdomain) return (zero_row, zero_column) else: - rhs = F_target.subs(targets_to_arrays(arrays['x'], targets)) + rhs = F_target.subs(targets_to_arrays(x, targets)) rhs = rhs.subs(self.time_mapper) - return as_tuple(Eq(arrays['y'], rhs, subdomain=eq.subdomain)) + return as_tuple(Eq(y, rhs, subdomain=eq.subdomain)) def _scale_non_bcs(self, matvecs, target=None): target = target or self.target + # TODO: make this a property of the class so don't need target as an arg vol = target.grid.symbolic_volume_cell return [ @@ -306,13 +319,13 @@ def _scale_non_bcs(self, matvecs, target=None): for m in matvecs ] - def _compute_scdiag(self, matvecs, arrays=None): + def _compute_scdiag(self, matvecs, col_target=None): """ """ - arrays = arrays or self.arrays + x = self.arrays[col_target or self.target]['x'] centres = { - centre_stencil(m.rhs, arrays['x'], as_coeff=True) + centre_stencil(m.rhs, x, as_coeff=True) for m in matvecs if not isinstance(m, EssentialBC) } # add comments @@ -329,8 +342,8 @@ def _scale_bcs(self, matvecs, scdiag): class SubMatrixBlock: - def __init__(self, name, matvecs, scdiag, - row_target, col_target, row_idx, col_idx, linear_idx): + def __init__(self, name, matvecs, scdiag, row_target, + col_target, row_idx, col_idx, linear_idx): self.name = name self.matvecs = matvecs self.scdiag = scdiag @@ -408,11 +421,10 @@ def _build_blocks(self, target_eqns): """ for i, row_target in enumerate(self.targets): eqns = target_eqns[row_target] - arrays = self.arrays[row_target] for j, col_target in enumerate(self.targets): matvecs = [ e for eq in eqns for e in - self._build_matvec_eq(eq, col_target, arrays) + self._build_matvec_eq(eq, col_target, row_target) ] matvecs = [m for m in matvecs if m is not None] # Sort to put EssentialBC first if any @@ -420,7 +432,7 @@ def _build_blocks(self, target_eqns): sorted(matvecs, key=lambda e: not isinstance(e, EssentialBC)) ) matvecs = self._scale_non_bcs(matvecs, row_target) - scdiag = self._compute_scdiag(matvecs, arrays) + scdiag = self._compute_scdiag(matvecs, col_target) matvecs = self._scale_bcs(matvecs, scdiag) name = f'J{i}{j}' @@ -492,7 +504,7 @@ def _build_equations(self): self._formrhs = rhs def _make_F_target(self, eq, F_target, targets): - arrays = self.arrays + arrays = self.arrays[self.target] volume = self.target.grid.symbolic_volume_cell if isinstance(eq, EssentialBC): # The initial guess satisfies the essential BCs, so this term is zero. @@ -511,9 +523,10 @@ def _make_F_target(self, eq, F_target, targets): return as_tuple(Eq(arrays['f'], rhs, subdomain=eq.subdomain)) def _make_b(self, eq, b): + b_arr = self.arrays[self.target]['b'] rhs = 0. if isinstance(eq, EssentialBC) else b.subs(self.time_mapper) rhs = rhs * self.target.grid.symbolic_volume_cell - return as_tuple(Eq(self.arrays['b'], rhs, subdomain=eq.subdomain)) + return as_tuple(Eq(b_arr, rhs, subdomain=eq.subdomain)) def _scale_bcs(self, eq, scdiag=None): """ @@ -616,7 +629,7 @@ def _make_initial_guess(self, eq): if isinstance(eq, EssentialBC): assert eq.lhs == self.target return Eq( - self.arrays['x'], eq.rhs, + self.arrays[self.target]['x'], eq.rhs, subdomain=eq.subdomain ) else: diff --git a/devito/types/array.py b/devito/types/array.py index 3c3cedfc88..b3367f44a9 100644 --- a/devito/types/array.py +++ b/devito/types/array.py @@ -538,7 +538,7 @@ class ComponentAccess(Expr, Pickable): _default_component_names = ('x', 'y', 'z', 'w') __rargs__ = ('arg',) - __rkwargs__ = ('index',) + __rkwargs__ = ('index', 'component_names') def __new__(cls, arg, index=0, component_names=None, **kwargs): if not arg.is_Indexed: @@ -579,6 +579,10 @@ def arg(self): def index(self): return self._index + @property + def component_names(self): + return self._component_names + @property def sindex(self): return self._component_names[self.index] diff --git a/examples/petsc/Poisson/01_poisson.py b/examples/petsc/Poisson/01_poisson.py index 4981566af6..6f24f74f64 100644 --- a/examples/petsc/Poisson/01_poisson.py +++ b/examples/petsc/Poisson/01_poisson.py @@ -66,6 +66,7 @@ def analytical(x, y): Ly = np.float64(1.) n_values = list(range(13, 174, 10)) +n_values = [9] dx = np.array([Lx/(n-1) for n in n_values]) errors = [] diff --git a/examples/petsc/cfd/01_navierstokes.py b/examples/petsc/cfd/01_navierstokes.py index 01e96e1afe..7797321045 100644 --- a/examples/petsc/cfd/01_navierstokes.py +++ b/examples/petsc/cfd/01_navierstokes.py @@ -1,8 +1,8 @@ import os import numpy as np -from devito import (Grid, TimeFunction, Function, Constant, Eq, - Operator, norm, SubDomain, switchconfig, configuration) +from devito import (Grid, TimeFunction, Constant, Eq, + Operator, SubDomain, switchconfig, configuration) from devito.symbolics import retrieve_functions, INT from devito.petsc import PETScSolve, EssentialBC @@ -232,11 +232,9 @@ def neumann_right(eq, subdomain): u1 = TimeFunction(name='u1', grid=grid, space_order=2, dtype=np.float64) v1 = TimeFunction(name='v1', grid=grid, space_order=2, dtype=np.float64) -pn1 = Function(name='pn1', grid=grid, space_order=2, dtype=np.float64) +pn1 = TimeFunction(name='pn1', grid=grid, space_order=2, dtype=np.float64) -pn1.data[:] = 0. - -eq_pn1 = Eq(pn1.laplace, rho*(1./dt*(u1.forward.dxc+v1.forward.dyc)), +eq_pn1 = Eq(pn1.forward.laplace, rho*(1./dt*(u1.forward.dxc+v1.forward.dyc)), subdomain=grid.interior) @@ -244,21 +242,21 @@ def neumann_right(eq, subdomain): bc_pn1 += [neumann_bottom(eq_pn1, sub2)] bc_pn1 += [neumann_left(eq_pn1, sub3)] bc_pn1 += [neumann_right(eq_pn1, sub4)] -bc_pn1 += [EssentialBC(pn1, 0., subdomain=sub5)] +bc_pn1 += [EssentialBC(pn1.forward, 0., subdomain=sub5)] bc_pn1 += [neumann_right(neumann_bottom(eq_pn1, sub6), sub6)] bc_pn1 += [neumann_left(neumann_top(eq_pn1, sub7), sub7)] bc_pn1 += [neumann_right(neumann_top(eq_pn1, sub8), sub8)] -eqn_p = PETScSolve([eq_pn1]+bc_pn1, pn1) +eqn_p = PETScSolve([eq_pn1]+bc_pn1, pn1.forward) -eq_u1 = Eq(u1.dt + u1*u1.dxc + v1*u1.dyc, nu*u1.laplace) -eq_v1 = Eq(v1.dt + u1*v1.dxc + v1*v1.dyc, nu*v1.laplace) +eq_u1 = Eq(u1.dt + u1*u1.dxc + v1*u1.dyc, nu*u1.laplace, subdomain=grid.interior) +eq_v1 = Eq(v1.dt + u1*v1.dxc + v1*v1.dyc, nu*v1.laplace, subdomain=grid.interior) -update_u = Eq(u1.forward, u1.forward - (dt/rho)*(pn1.dxc), +update_u = Eq(u1.forward, u1.forward - (dt/rho)*(pn1.forward.dxc), subdomain=grid.interior) -update_v = Eq(v1.forward, v1.forward - (dt/rho)*(pn1.dyc), +update_v = Eq(v1.forward, v1.forward - (dt/rho)*(pn1.forward.dyc), subdomain=grid.interior) # TODO: Can drop due to initial guess CB @@ -296,12 +294,6 @@ def neumann_right(eq, subdomain): op = Operator(exprs) op.apply(time_m=0, time_M=ns-1, dt=dt) -u1_norm = norm(u1) -v1_norm = norm(v1) -p1_norm = norm(pn1) - - -# TODO: change these norm checks to array checks (use paper) -assert np.isclose(u1_norm, 13.966067703420883, atol=0, rtol=1e-7) -assert np.isclose(v1_norm, 7.9575677674738285, atol=0, rtol=1e-7) -assert np.isclose(p1_norm, 36.46263134701362, atol=0, rtol=1e-7) +# Pressure norm check +tol = 1e-3 +assert np.sum((pn1.data[0]-pn1.data[1])**2/np.maximum(pn1.data[0]**2, 1e-10)) < tol diff --git a/tests/test_petsc.py b/tests/test_petsc.py index d0f4bdd087..e0707d7b1e 100644 --- a/tests/test_petsc.py +++ b/tests/test_petsc.py @@ -1,3 +1,5 @@ +import pytest + import numpy as np import os @@ -615,19 +617,38 @@ class TestCoupledLinear: # TODO: Add more comprehensive tests for fully coupled problems. # TODO: Add subdomain tests, time loop, multiple coupled etc. + @pytest.mark.parametrize('eq1, eq2, so', [ + ('Eq(e.laplace, f)', 'Eq(g.laplace, h)', '2'), + ('Eq(e.laplace, f)', 'Eq(g.laplace, h)', '4'), + ('Eq(e.laplace, f)', 'Eq(g.laplace, h)', '6'), + ('Eq(e.laplace, f + 5)', 'Eq(g.laplace, h + 5)', '2'), + ('Eq(e.laplace, f + 5)', 'Eq(g.laplace, h + 5)', '4'), + ('Eq(e.laplace, f + 5)', 'Eq(g.laplace, h + 5)', '6'), + ('Eq(e.dx, e + 2*f)', 'Eq(g.dx, g + 2*h)', '2'), + ('Eq(e.dx, e + 2*f)', 'Eq(g.dx, g + 2*h)', '4'), + ('Eq(e.dx, e + 2*f)', 'Eq(g.dx, g + 2*h)', '6'), + ('Eq(f.dx, e.dx + e + e.laplace)', 'Eq(h.dx, g.dx + g + g.laplace)', '2'), + ('Eq(f.dx, e.dx + e + e.laplace)', 'Eq(h.dx, g.dx + g + g.laplace)', '4'), + ('Eq(f.dx, e.dx + e + e.laplace)', 'Eq(h.dx, g.dx + g + g.laplace)', '6'), + ]) @skipif('petsc') - def test_coupled_vs_non_coupled(self): + def test_coupled_vs_non_coupled(self, eq1, eq2, so): + """ + Test that solving multiple **uncoupled** equations separately + vs. together with `PETScSolve` yields the same result. + This test is non time-dependent. + """ grid = Grid(shape=(11, 11), dtype=np.float64) - functions = [Function(name=n, grid=grid, space_order=2) + functions = [Function(name=n, grid=grid, space_order=eval(so)) for n in ['e', 'f', 'g', 'h']] e, f, g, h = functions f.data[:] = 5. h.data[:] = 5. - eq1 = Eq(e.laplace, f) - eq2 = Eq(g.laplace, h) + eq1 = eval(eq1) + eq2 = eval(eq2) # Non-coupled petsc1 = PETScSolve(eq1, target=e) @@ -645,8 +666,6 @@ def test_coupled_vs_non_coupled(self): g.data[:] = 0 # Coupled - # TODO: Need more friendly API for coupled - just - # using a dict for now petsc3 = PETScSolve({e: [eq1], g: [eq2]}) with switchconfig(language='petsc'): @@ -658,8 +677,8 @@ def test_coupled_vs_non_coupled(self): print('enorm1:', enorm1) print('enorm2:', enorm2) - assert np.isclose(enorm1, enorm2, rtol=1e-16) - assert np.isclose(gnorm1, gnorm2, rtol=1e-16) + assert np.isclose(enorm1, enorm2, atol=1e-14) + assert np.isclose(gnorm1, gnorm2, atol=1e-14) callbacks1 = [meta_call.root for meta_call in op1._func_table.values()] callbacks2 = [meta_call.root for meta_call in op2._func_table.values()] @@ -724,49 +743,33 @@ def test_coupled_structs(self): assert 'struct Field0\n{' not in ccode assert 'struct Field0\n{' in hcode + @pytest.mark.parametrize('n_fields', [2, 3, 4, 5, 6]) @skipif('petsc') - def test_coupled_frees(self): + def test_coupled_frees(self, n_fields): grid = Grid(shape=(11, 11), dtype=np.float64) - functions = [Function(name=n, grid=grid, space_order=2) - for n in ['e', 'f', 'g', 'h']] - e, f, g, h = functions - - eq1 = Eq(e.laplace, h) - eq2 = Eq(f.laplace, h) - eq3 = Eq(g.laplace, h) + functions = [Function(name=f'u{i}', grid=grid, space_order=2) + for i in range(n_fields + 1)] + *solved_funcs, h = functions - petsc1 = PETScSolve({e: [eq1], f: [eq2]}) - petsc2 = PETScSolve({e: [eq1], f: [eq2], g: [eq3]}) + equations = [Eq(func.laplace, h) for func in solved_funcs] + petsc = PETScSolve({func: [eq] for func, eq in zip(solved_funcs, equations)}) with switchconfig(language='petsc'): - op1 = Operator(petsc1, opt='noop') - op2 = Operator(petsc2, opt='noop') + op = Operator(petsc, opt='noop') + + frees = op.body.frees - frees1 = op1.body.frees - frees2 = op2.body.frees - - # Check solver with two fields - # IS destroys - assert str(frees1[0]) == 'PetscCall(ISDestroy(&(fields0[0])));' - assert str(frees1[1]) == 'PetscCall(ISDestroy(&(fields0[1])));' - assert str(frees1[2]) == 'PetscCall(PetscFree(fields0));' - # Sub DM destroys - assert str(frees1[3]) == 'PetscCall(DMDestroy(&(subdms0[0])));' - assert str(frees1[4]) == 'PetscCall(DMDestroy(&(subdms0[1])));' - assert str(frees1[5]) == 'PetscCall(PetscFree(subdms0));' - - # Check solver with three fields - # IS destroys - assert str(frees2[0]) == 'PetscCall(ISDestroy(&(fields0[0])));' - assert str(frees2[1]) == 'PetscCall(ISDestroy(&(fields0[1])));' - assert str(frees2[2]) == 'PetscCall(ISDestroy(&(fields0[2])));' - assert str(frees2[3]) == 'PetscCall(PetscFree(fields0));' - # Sub DM destroys - assert str(frees2[4]) == 'PetscCall(DMDestroy(&(subdms0[0])));' - assert str(frees2[5]) == 'PetscCall(DMDestroy(&(subdms0[1])));' - assert str(frees2[6]) == 'PetscCall(DMDestroy(&(subdms0[2])));' - assert str(frees2[7]) == 'PetscCall(PetscFree(subdms0));' + # IS Destroy calls + for i in range(n_fields): + assert str(frees[i]) == f'PetscCall(ISDestroy(&(fields0[{i}])));' + assert str(frees[n_fields]) == 'PetscCall(PetscFree(fields0));' + + # DM Destroy calls + for i in range(n_fields): + assert str(frees[n_fields + 1 + i]) == \ + f'PetscCall(DMDestroy(&(subdms0[{i}])));' + assert str(frees[n_fields*2 + 1]) == 'PetscCall(PetscFree(subdms0));' @skipif('petsc') def test_dmda_dofs(self): @@ -803,7 +806,7 @@ def test_dmda_dofs(self): in str(op3) @skipif('petsc') - def test_submatrices(self): + def test_mixed_jacobian(self): grid = Grid(shape=(11, 11), dtype=np.float64) functions = [Function(name=n, grid=grid, space_order=2) @@ -849,6 +852,169 @@ def test_submatrices(self): assert j10.col_target == e assert j11.col_target == g + + # TODO: FIX THIS LOGIC - THE OUTPUT FOR J01_MATVEC IS Eq(y_e(x, y), -h_x*h_y*x_e(x, y)) BUT IT SHOULD + # BE Eq(y_e(x, y), -h_x*h_y*x_G(x, y)) + + # @pytest.mark.parametrize('eq1, eq2, so, j01_matvec, j10_matvec', [ + # ('Eq(-e.laplace, g)', 'Eq(-g.laplace, e)', '2', 'h_x*h_y*(-2.0/h_y**2 - 2.0/h_x**2)'), + # ]) + # @skipif('petsc') + # def test_coupling(self, eq1, eq2, so, scale): + # """ + # Test linear coupling between fields, where the off-diagonal + # Jacobian submatrices are nonzero. + # """ + # grid = Grid(shape=(9,9), dtype=np.float64) + + # functions = [Function(name=n, grid=grid, space_order=eval(so)) + # for n in ['e', 'f', 'g', 'h']] + # e, f, g, h = functions + + # eq1 = eval(eq1) + # eq2 = eval(eq2) + + # petsc = PETScSolve({e: [eq1], g: [eq2]}) + + # jacobian = petsc[0].rhs.fielddata.jacobian + + # j01 = jacobian.get_submatrix(0, 1) + # j10 = jacobian.get_submatrix(1, 0) + + # assert len(j01.matvecs) == 1 + # assert len(j10.matvecs) == 1 + + # assert j01.col_target == g + # assert j10.col_target == e + + # assert str(j01.matvecs[0]) == j01_matvec + # assert str(j10.matvecs[0]) == j10_matvec + + @pytest.mark.parametrize('eq1, eq2, so, scale', [ + ('Eq(e.laplace, f)', 'Eq(g.laplace, h)', '2', '-2.0*h_x/h_x**2'), + ('Eq(e.laplace, f)', 'Eq(g.laplace, h)', '4', '-2.5*h_x/h_x**2'), + ('Eq(e.laplace + e, f)', 'Eq(g.laplace + g, h)', '2', 'h_x*(1 - 2.0/h_x**2)'), + ('Eq(e.laplace + e, f)', 'Eq(g.laplace + g, h)', '4', 'h_x*(1 - 2.5/h_x**2)'), + ('Eq(e.laplace + 5.*e, f)', 'Eq(g.laplace + 5.*g, h)', '2', 'h_x*(5.0 - 2.0/h_x**2)'), + ('Eq(e.laplace + 5.*e, f)', 'Eq(g.laplace + 5.*g, h)', '4', 'h_x*(5.0 - 2.5/h_x**2)'), + ('Eq(e.dx + e + e.laplace, f)', 'Eq(g.dx + g + g.laplace, h.dx)', '2', 'h_x*(1 + 1/h_x - 2.0/h_x**2)'), + ('Eq(e.dx + e + e.laplace, f)', 'Eq(g.dx + g + g.laplace, h.dx)', '4', 'h_x*(1 - 2.5/h_x**2)'), + ('Eq(2.*e.laplace + e, f)', 'Eq(2*g.laplace + g, h)', '2', 'h_x*(1 - 4.0/h_x**2)'), + ('Eq(2.*e.laplace + e, f)', 'Eq(2*g.laplace + g, h)', '4', 'h_x*(1 - 5.0/h_x**2)'), + ]) + @skipif('petsc') + def test_jacobian_scaling_1D(self, eq1, eq2, so, scale): + """ + Test the computation of diagonal scaling in a 1D Jacobian system. + + This scaling would be applied to the boundary rows of the matrix + if essential boundary conditions were enforced in the solver. + Its purpose is to reduce the condition number of the matrix. + """ + grid = Grid(shape=(9,), dtype=np.float64) + + functions = [Function(name=n, grid=grid, space_order=eval(so)) + for n in ['e', 'f', 'g', 'h']] + e, f, g, h = functions + + eq1 = eval(eq1) + eq2 = eval(eq2) + + petsc = PETScSolve({e: [eq1], g: [eq2]}) + + jacobian = petsc[0].rhs.fielddata.jacobian + + j00 = jacobian.get_submatrix(0, 0) + j11 = jacobian.get_submatrix(1, 1) + + assert str(j00.scdiag) == scale + assert str(j11.scdiag) == scale + + @pytest.mark.parametrize('eq1, eq2, so, scale', [ + ('Eq(e.laplace, f)', 'Eq(g.laplace, h)', '2', 'h_x*h_y*(-2.0/h_y**2 - 2.0/h_x**2)'), + ('Eq(e.laplace, f)', 'Eq(g.laplace, h)', '4', 'h_x*h_y*(-2.5/h_y**2 - 2.5/h_x**2)'), + ('Eq(e.laplace + e, f)', 'Eq(g.laplace + g, h)', '2', 'h_x*h_y*(1 - 2.0/h_y**2 - 2.0/h_x**2)'), + ('Eq(e.laplace + e, f)', 'Eq(g.laplace + g, h)', '4', 'h_x*h_y*(1 - 2.5/h_y**2 - 2.5/h_x**2)'), + ('Eq(e.laplace + 5.*e, f)', 'Eq(g.laplace + 5.*g, h)', '2', 'h_x*h_y*(5.0 - 2.0/h_y**2 - 2.0/h_x**2)'), + ('Eq(e.laplace + 5.*e, f)', 'Eq(g.laplace + 5.*g, h)', '4', 'h_x*h_y*(5.0 - 2.5/h_y**2 - 2.5/h_x**2)'), + ('Eq(e.dx + e.dy + e + e.laplace, f)', 'Eq(g.dx + g.dy + g + g.laplace, h)', '2', 'h_x*h_y*(1 + 1/h_y - 2.0/h_y**2 + 1/h_x - 2.0/h_x**2)'), + ('Eq(e.dx + e.dy + e + e.laplace, f)', 'Eq(g.dx + g.dy + g + g.laplace, h)', '4', 'h_x*h_y*(1 - 2.5/h_y**2 - 2.5/h_x**2)'), + ('Eq(2.*e.laplace + e, f)', 'Eq(2*g.laplace + g, h)', '2', 'h_x*h_y*(1 - 4.0/h_y**2 - 4.0/h_x**2)'), + ('Eq(2.*e.laplace + e, f)', 'Eq(2*g.laplace + g, h)', '4', 'h_x*h_y*(1 - 5.0/h_y**2 - 5.0/h_x**2)'), + + ]) + @skipif('petsc') + def test_jacobian_scaling_2D(self, eq1, eq2, so, scale): + """ + Test the computation of diagonal scaling in a 2D Jacobian system. + + This scaling would be applied to the boundary rows of the matrix + if essential boundary conditions were enforced in the solver. + Its purpose is to reduce the condition number of the matrix. + """ + grid = Grid(shape=(9,9), dtype=np.float64) + + functions = [Function(name=n, grid=grid, space_order=eval(so)) + for n in ['e', 'f', 'g', 'h']] + e, f, g, h = functions + + eq1 = eval(eq1) + eq2 = eval(eq2) + + petsc = PETScSolve({e: [eq1], g: [eq2]}) + + jacobian = petsc[0].rhs.fielddata.jacobian + + j00 = jacobian.get_submatrix(0, 0) + j11 = jacobian.get_submatrix(1, 1) + + assert str(j00.scdiag) == scale + assert str(j11.scdiag) == scale + + @pytest.mark.parametrize('eq1, eq2, so, scale', [ + ('Eq(e.laplace, f)', 'Eq(g.laplace, h)', '2', 'h_x*h_y*h_z*(-2.0/h_z**2 - 2.0/h_y**2 - 2.0/h_x**2)'), + ('Eq(e.laplace, f)', 'Eq(g.laplace, h)', '4', 'h_x*h_y*h_z*(-2.5/h_z**2 - 2.5/h_y**2 - 2.5/h_x**2)'), + ('Eq(e.laplace + e, f)', 'Eq(g.laplace + g, h)', '2', 'h_x*h_y*h_z*(1 - 2.0/h_z**2 - 2.0/h_y**2 - 2.0/h_x**2)'), + ('Eq(e.laplace + e, f)', 'Eq(g.laplace + g, h)', '4', 'h_x*h_y*h_z*(1 - 2.5/h_z**2 - 2.5/h_y**2 - 2.5/h_x**2)'), + ('Eq(e.laplace + 5.*e, f)', 'Eq(g.laplace + 5.*g, h)', '2', 'h_x*h_y*h_z*(5.0 - 2.0/h_z**2 - 2.0/h_y**2 - 2.0/h_x**2)'), + ('Eq(e.laplace + 5.*e, f)', 'Eq(g.laplace + 5.*g, h)', '4', 'h_x*h_y*h_z*(5.0 - 2.5/h_z**2 - 2.5/h_y**2 - 2.5/h_x**2)'), + ('Eq(e.dx + e.dy + e.dz + e + e.laplace, f)', 'Eq(g.dx + g.dy + g.dz + g + g.laplace, h)', '2', 'h_x*h_y*h_z*(1 + 1/h_z - 2.0/h_z**2 + 1/h_y - 2.0/h_y**2 + 1/h_x - 2.0/h_x**2)'), + ('Eq(e.dx + e.dy + e.dz + e + e.laplace, f)', 'Eq(g.dx + g.dy + g.dz + g + g.laplace, h)', '4', 'h_x*h_y*h_z*(1 - 2.5/h_z**2 - 2.5/h_y**2 - 2.5/h_x**2)'), + ('Eq(2.*e.laplace + e, f)', 'Eq(2*g.laplace + g, h)', '2', 'h_x*h_y*h_z*(1 - 4.0/h_z**2 - 4.0/h_y**2 - 4.0/h_x**2)'), + ('Eq(2.*e.laplace + e, f)', 'Eq(2*g.laplace + g, h)', '4', 'h_x*h_y*h_z*(1 - 5.0/h_z**2 - 5.0/h_y**2 - 5.0/h_x**2)'), + + ]) + @skipif('petsc') + def test_jacobian_scaling_3D(self, eq1, eq2, so, scale): + """ + Test the computation of diagonal scaling in a 3D Jacobian system. + + This scaling would be applied to the boundary rows of the matrix + if essential boundary conditions were enforced in the solver. + Its purpose is to reduce the condition number of the matrix. + """ + grid = Grid(shape=(9,9,9), dtype=np.float64) + + functions = [Function(name=n, grid=grid, space_order=eval(so)) + for n in ['e', 'f', 'g', 'h']] + e, f, g, h = functions + + eq1 = eval(eq1) + eq2 = eval(eq2) + + petsc = PETScSolve({e: [eq1], g: [eq2]}) + + jacobian = petsc[0].rhs.fielddata.jacobian + + j00 = jacobian.get_submatrix(0, 0) + j11 = jacobian.get_submatrix(1, 1) + + assert str(j00.scdiag) == scale + assert str(j11.scdiag) == scale + + # test coupled residual callback - check the .dot for each field etc + + @skipif('petsc') def test_residual_bundle(self): grid = Grid(shape=(11, 11), dtype=np.float64) From 5c1b686938507d73c43392316fcc15890d2e5133 Mon Sep 17 00:00:00 2001 From: ZoeLeibowitz Date: Sun, 8 Jun 2025 12:46:10 +0100 Subject: [PATCH 41/84] tests: Add more petsc tests for linear coupled --- .github/workflows/examples-mpi.yml | 4 +- .github/workflows/examples.yml | 4 +- .github/workflows/flake8.yml | 4 +- .github/workflows/pytest-core-mpi.yml | 4 +- .github/workflows/pytest-core-nompi.yml | 4 +- .github/workflows/pytest-petsc.yml | 4 +- .github/workflows/tutorials.yml | 4 +- devito/passes/iet/engine.py | 9 +- devito/petsc/iet/routines.py | 24 +-- devito/petsc/solve.py | 11 +- devito/petsc/types/equation.py | 10 +- devito/petsc/types/types.py | 10 +- devito/tools/utils.py | 2 +- examples/petsc/Poisson/01_poisson.py | 3 - examples/petsc/Poisson/02_laplace.py | 2 - examples/petsc/petsc_test.py | 2 +- examples/petsc/random/02_biharmonic.py | 4 +- tests/test_petsc.py | 213 ++++++++++++++++-------- 18 files changed, 190 insertions(+), 128 deletions(-) diff --git a/.github/workflows/examples-mpi.yml b/.github/workflows/examples-mpi.yml index e44a03410d..47a8629c05 100644 --- a/.github/workflows/examples-mpi.yml +++ b/.github/workflows/examples-mpi.yml @@ -17,11 +17,11 @@ on: push: branches: - main - - master + - petsc pull_request: branches: - main - - master + - petsc jobs: build: diff --git a/.github/workflows/examples.yml b/.github/workflows/examples.yml index 9b6d5a9ade..4ca681131f 100644 --- a/.github/workflows/examples.yml +++ b/.github/workflows/examples.yml @@ -10,11 +10,11 @@ on: push: branches: - main - - master + - petsc pull_request: branches: - main - - master + - petsc jobs: tutorials: diff --git a/.github/workflows/flake8.yml b/.github/workflows/flake8.yml index 40b57d7706..e207557f10 100644 --- a/.github/workflows/flake8.yml +++ b/.github/workflows/flake8.yml @@ -10,11 +10,11 @@ on: push: branches: - main - - master + - petsc pull_request: branches: - main - - master + - petsc jobs: flake8: diff --git a/.github/workflows/pytest-core-mpi.yml b/.github/workflows/pytest-core-mpi.yml index fbee3367f1..07ca0534d6 100644 --- a/.github/workflows/pytest-core-mpi.yml +++ b/.github/workflows/pytest-core-mpi.yml @@ -10,11 +10,11 @@ on: push: branches: - main - - master + - petsc pull_request: branches: - main - - master + - petsc jobs: test-mpi-basic: diff --git a/.github/workflows/pytest-core-nompi.yml b/.github/workflows/pytest-core-nompi.yml index 23b5b20344..b00115ba7f 100644 --- a/.github/workflows/pytest-core-nompi.yml +++ b/.github/workflows/pytest-core-nompi.yml @@ -10,11 +10,11 @@ on: push: branches: - main - - master + - petsc pull_request: branches: - main - - master + - petsc jobs: pytest: diff --git a/.github/workflows/pytest-petsc.yml b/.github/workflows/pytest-petsc.yml index 9c035a86e7..2bf55b5049 100644 --- a/.github/workflows/pytest-petsc.yml +++ b/.github/workflows/pytest-petsc.yml @@ -10,11 +10,11 @@ on: push: branches: - main - - master + - petsc pull_request: branches: - main - - master + - petsc jobs: pytest: diff --git a/.github/workflows/tutorials.yml b/.github/workflows/tutorials.yml index c6e9c2adba..75b7ef3982 100644 --- a/.github/workflows/tutorials.yml +++ b/.github/workflows/tutorials.yml @@ -10,11 +10,11 @@ on: push: branches: - main - - master + - petsc pull_request: branches: - main - - master + - petsc jobs: tutorials: diff --git a/devito/passes/iet/engine.py b/devito/passes/iet/engine.py index ad0ebefd1e..f141f436d3 100644 --- a/devito/passes/iet/engine.py +++ b/devito/passes/iet/engine.py @@ -1,6 +1,7 @@ from collections import OrderedDict, defaultdict from functools import partial, singledispatch, wraps +from devito.finite_differences.differentiable import Differentiable from devito.ir.iet import (Call, ExprStmt, Iteration, SyncSpot, AsyncCallable, FindNodes, FindSymbols, MapNodes, MetaCall, Transformer, EntryFunction, FixedArgsCallable, Uxreplace, @@ -10,13 +11,13 @@ from devito.passes import needs_transfer from devito.symbolics import FieldFromComposite, FieldFromPointer from devito.tools import DAG, as_tuple, filter_ordered, sorted_priority, timed_pass -from devito.types import (Bundle, CompositeObject, Lock, IncrDimension, +from devito.types import (Array, Bundle, CompositeObject, Lock, IncrDimension, ModuloDimension, Indirection, Pointer, SharedData, ThreadArray, Temp, NPThreads, NThreadsBase, Wildcard) -from devito.types.array import ArrayBasic from devito.types.args import ArgProvider from devito.types.dense import DiscreteFunction from devito.types.dimension import AbstractIncrDimension, BlockDimension +from devito.types.array import ArrayBasic __all__ = ['Graph', 'iet_pass', 'iet_visit'] @@ -414,7 +415,8 @@ def abstract_objects(objects0, sregistry=None): # Precedence rules make it possible to reconstruct objects that depend on # higher priority objects - keys = [Bundle, ArrayBasic, DiscreteFunction, AbstractIncrDimension, BlockDimension] + keys = [Bundle, Array, Differentiable, DiscreteFunction, + AbstractIncrDimension, BlockDimension] priority = {k: i for i, k in enumerate(keys, start=1)} objects = sorted_priority(objects, priority) @@ -449,6 +451,7 @@ def _(i, mapper, sregistry): }) +@abstract_object.register(Array) @abstract_object.register(ArrayBasic) def _(i, mapper, sregistry): if isinstance(i, Lock): diff --git a/devito/petsc/iet/routines.py b/devito/petsc/iet/routines.py index 2fbdafe1ef..4e48a571f2 100644 --- a/devito/petsc/iet/routines.py +++ b/devito/petsc/iet/routines.py @@ -688,15 +688,11 @@ def zero_memory(self): return False def _make_core(self): - injectsolve = self.injectsolve - all_fielddata = injectsolve.expr.rhs.fielddata - - for sm in all_fielddata.jacobian.nonzero_submatrices: - # arrays = all_fielddata.arrays[sm.row_target] + for sm in self.fielddata.jacobian.nonzero_submatrices: self._make_matvec(sm, prefix=f'{sm.name}_MatMult') self._make_whole_matvec() - self._make_whole_formfunc(all_fielddata) + self._make_whole_formfunc() self._make_user_struct_callback() self._create_submatrices() self._efuncs['PopulateMatContext'] = self.objs['dummyefunc'] @@ -764,15 +760,15 @@ def _whole_matvec_body(self): retstmt=(Call('PetscFunctionReturn', arguments=[0]),) ) - def _make_whole_formfunc(self, fielddata): - formfuncs = fielddata.residual.formfuncs + def _make_whole_formfunc(self): + formfuncs = self.fielddata.residual.formfuncs # Compile formfunc `eqns` into an IET via recursive compilation irs_formfunc, _ = self.rcompile( formfuncs, options={'mpi': False}, sregistry=self.sregistry, concretize_mapper=self.concretize_mapper ) - body_formfunc = self._whole_formfunc_body(List(body=irs_formfunc.uiet.body), - fielddata) + body_formfunc = self._whole_formfunc_body(List(body=irs_formfunc.uiet.body)) + objs = self.objs cb = PETScCallable( self.sregistry.make_name(prefix='WholeFormFunc'), @@ -783,7 +779,7 @@ def _make_whole_formfunc(self, fielddata): self._main_formfunc_callback = cb self._efuncs[cb.name] = cb - def _whole_formfunc_body(self, body, fielddata): + def _whole_formfunc_body(self, body): linsolve_expr = self.injectsolve.expr.rhs objs = self.objs sobjs = self.solver_objs @@ -796,7 +792,7 @@ def _whole_formfunc_body(self, body, fielddata): fields = self._dummy_fields(body) self._struct_params.extend(fields) - # Process body for residual callback, including generating bundles etc + # Process body with bundles for residual callback bundles = sobjs['bundles'] fbundle = bundles['f'] xbundle = bundles['x'] @@ -1143,7 +1139,6 @@ def _extend_build(self, base_dict): class CoupledObjectBuilder(BaseObjectBuilder): def _extend_build(self, base_dict): - injectsolve = self.injectsolve sreg = self.sregistry objs = self.objs targets = self.fielddata.targets @@ -1164,8 +1159,7 @@ def _extend_build(self, base_dict): dim_labels[i]: PetscInt(dim_labels[i]) for i in range(space_dims) }) - jacobian = injectsolve.expr.rhs.fielddata.jacobian - submatrices = jacobian.nonzero_submatrices + submatrices = self.fielddata.jacobian.nonzero_submatrices base_dict['jacctx'] = JacobianStruct( name=sreg.make_name(prefix=objs['ljacctx'].name), diff --git a/devito/petsc/solve.py b/devito/petsc/solve.py index 8adb9cda1e..61d0c187e6 100644 --- a/devito/petsc/solve.py +++ b/devito/petsc/solve.py @@ -27,12 +27,12 @@ def PETScSolve(target_eqns, target=None, solver_parameters=None): target_eqns : Eq or list of Eq, or dict of Function-like -> Eq or list of Eq The targets and symbolic equations defining the system to be solved. - - **Single-field problem**: + - Single-field problem: Pass a single Eq or list of Eq, and specify `target` separately: PETScSolve(Eq1, target) PETScSolve([Eq1, Eq2], target) - - **Multi-field (mixed) problem**: + - Multi-field (mixed) problem: Pass a dictionary mapping each target field to its Eq(s): PETScSolve({u: Eq1, v: Eq2}) PETScSolve({u: [Eq1, Eq2], v: [Eq3, Eq4]}) @@ -47,8 +47,8 @@ def PETScSolve(target_eqns, target=None, solver_parameters=None): Returns ------- Eq - A symbolic equation that wraps the system, solver metadata, - and boundary conditions. This can be passed directly to a Devito Operator. + A symbolic equation that wraps the linear solver. + This can be passed directly to a Devito Operator. """ if target is not None: return InjectSolve(solver_parameters, {target: target_eqns}).build_eq() @@ -143,9 +143,6 @@ def linear_solve_args(self): return coupled_targets[0], tuple(funcs), all_data - # def generate_arrays_combined(self, *targets): - # return {target: self.generate_arrays(target) for target in targets} - def generate_time_mapper(funcs): """ diff --git a/devito/petsc/types/equation.py b/devito/petsc/types/equation.py index 6d476c12f7..e819b48a22 100644 --- a/devito/petsc/types/equation.py +++ b/devito/petsc/types/equation.py @@ -27,8 +27,9 @@ class ZeroRow(EssentialBC): Equation used to zero all entries, except the diagonal, of a row in the Jacobian. - Note: - This is only used interally by the compiler, not by users. + Warnings + -------- + Created and managed directly by Devito, not by users. """ pass @@ -37,7 +38,8 @@ class ZeroColumn(EssentialBC): """ Equation used to zero the column of the Jacobian. - Note: - This is only used interally by the compiler, not by users. + Warnings + -------- + Created and managed directly by Devito, not by users. """ pass diff --git a/devito/petsc/types/types.py b/devito/petsc/types/types.py index 21433b5944..44d43d2f97 100644 --- a/devito/petsc/types/types.py +++ b/devito/petsc/types/types.py @@ -236,7 +236,7 @@ def targets(self): class Jacobian: """ - Represents a Jacobian matrix in a matrix-free form. + Represents a Jacobian matrix. This Jacobian is defined implicitly via matrix-vector products derived from the symbolic equations provided in `matvecs`. @@ -311,7 +311,6 @@ def _make_matvec(self, eq, F_target, targets, col_target, row_target): def _scale_non_bcs(self, matvecs, target=None): target = target or self.target - # TODO: make this a property of the class so don't need target as an arg vol = target.grid.symbolic_volume_cell return [ @@ -328,7 +327,6 @@ def _compute_scdiag(self, matvecs, col_target=None): centre_stencil(m.rhs, x, as_coeff=True) for m in matvecs if not isinstance(m, EssentialBC) } - # add comments return centres.pop() if len(centres) == 1 else 1.0 def _scale_bcs(self, matvecs, scdiag): @@ -370,7 +368,7 @@ class MixedJacobian(Jacobian): as monolithic or block-structured in PETSc, but sub-block callbacks are generated in both cases. - Assumes a **linear** problem, so this Jacobian corresponds to a + Assumes a linear problem, so this Jacobian corresponds to a coefficient matrix and does not require differentiation. # TODO: pcfieldsplit support for each block @@ -407,7 +405,8 @@ def nonzero_submatrices(self): @property def target_scaler_mapper(self): """ - Map each row_target index to its diagonal submatrix's scaler. + Map each row target to the scdiag of its corresponding + diagonal subblock. """ mapper = {} for sm in self.submatrices: @@ -601,7 +600,6 @@ def _build_function_eq(self, eq, target): class InitialGuess: """ Enforce initial guess to satisfy essential BCs. - # TODO: Extend this to "coupled". """ def __init__(self, target, eqns, arrays): self.target = target diff --git a/devito/tools/utils.py b/devito/tools/utils.py index 1285ffabec..0a28de16a8 100644 --- a/devito/tools/utils.py +++ b/devito/tools/utils.py @@ -337,7 +337,7 @@ def sorted_priority(items, priority): """ def key(i): - for cls in sorted(priority, key=priority.get): + for cls in sorted(priority, key=priority.get, reverse=True): if isinstance(i, cls): v = priority[cls] break diff --git a/examples/petsc/Poisson/01_poisson.py b/examples/petsc/Poisson/01_poisson.py index 6f24f74f64..7ed32e8bbd 100644 --- a/examples/petsc/Poisson/01_poisson.py +++ b/examples/petsc/Poisson/01_poisson.py @@ -16,8 +16,6 @@ # Subdomains to implement BCs -# NOTE: For essential BCs, we ensure the SubDomains do not overlap - class SubTop(SubDomain): name = 'subtop' @@ -66,7 +64,6 @@ def analytical(x, y): Ly = np.float64(1.) n_values = list(range(13, 174, 10)) -n_values = [9] dx = np.array([Lx/(n-1) for n in n_values]) errors = [] diff --git a/examples/petsc/Poisson/02_laplace.py b/examples/petsc/Poisson/02_laplace.py index c84855137f..9df68f9ab9 100644 --- a/examples/petsc/Poisson/02_laplace.py +++ b/examples/petsc/Poisson/02_laplace.py @@ -23,8 +23,6 @@ # Subdomains to implement BCs -# NOTE: For essential BCs, we ensure the SubDomains do not overlap - class SubTop(SubDomain): name = 'subtop' diff --git a/examples/petsc/petsc_test.py b/examples/petsc/petsc_test.py index 12d245480e..5d93669d5f 100644 --- a/examples/petsc/petsc_test.py +++ b/examples/petsc/petsc_test.py @@ -20,7 +20,7 @@ v.data[:] = 5.0 -eq = Eq(0., u.laplace, subdomain=grid.interior) +eq = Eq(v, u.laplace, subdomain=grid.interior) petsc = PETScSolve([eq], u) diff --git a/examples/petsc/random/02_biharmonic.py b/examples/petsc/random/02_biharmonic.py index 214e38afda..f08ffc07de 100644 --- a/examples/petsc/random/02_biharmonic.py +++ b/examples/petsc/random/02_biharmonic.py @@ -85,7 +85,6 @@ def f_fcn(x, y): Ly = np.float64(1.) n_values = [33, 53, 73, 93, 113] -n_values = [9] dx = np.array([Lx/(n-1) for n in n_values]) u_errors = [] @@ -129,8 +128,7 @@ def f_fcn(x, y): with switchconfig(language='petsc'): op = Operator(petsc) - # op.apply() - print(op.ccode) + op.apply() u_exact.data[:] = u_exact_fcn(X, Y) lap_u.data[:] = lap_u_exact_fcn(X, Y) diff --git a/tests/test_petsc.py b/tests/test_petsc.py index e0707d7b1e..b1b9c3d187 100644 --- a/tests/test_petsc.py +++ b/tests/test_petsc.py @@ -852,63 +852,84 @@ def test_mixed_jacobian(self): assert j10.col_target == e assert j11.col_target == g + @pytest.mark.parametrize('eq1, eq2, j01_matvec, j10_matvec', [ + ('Eq(-e.laplace, g)', 'Eq(-g.laplace, e)', + 'Eq(y_e(x, y), -h_x*h_y*x_g(x, y))', + 'Eq(y_g(x, y), -h_x*h_y*x_e(x, y))'), + ('Eq(-e.laplace, 2.*g)', 'Eq(-g.laplace, 2.*e)', + 'Eq(y_e(x, y), -2.0*h_x*h_y*x_g(x, y))', + 'Eq(y_g(x, y), -2.0*h_x*h_y*x_e(x, y))'), + ('Eq(-e.laplace, g.dx)', 'Eq(-g.laplace, e.dx)', + 'Eq(y_e(x, y), -h_x*h_y*Derivative(x_g(x, y), x))', + 'Eq(y_g(x, y), -h_x*h_y*Derivative(x_e(x, y), x))'), + ('Eq(-e.laplace, g.dx + g)', 'Eq(-g.laplace, e.dx + e)', + 'Eq(y_e(x, y), h_x*h_y*(-x_g(x, y) - Derivative(x_g(x, y), x)))', + 'Eq(y_g(x, y), h_x*h_y*(-x_e(x, y) - Derivative(x_e(x, y), x)))'), + ('Eq(e, g.dx + g)', 'Eq(g, e.dx + e)', + 'Eq(y_e(x, y), h_x*h_y*(-x_g(x, y) - Derivative(x_g(x, y), x)))', + 'Eq(y_g(x, y), h_x*h_y*(-x_e(x, y) - Derivative(x_e(x, y), x)))'), + ('Eq(e, g.dx + g.dy)', 'Eq(g, e.dx + e.dy)', + 'Eq(y_e(x, y), h_x*h_y*(-Derivative(x_g(x, y), x) - Derivative(x_g(x, y), y)))', + 'Eq(y_g(x, y), h_x*h_y*(-Derivative(x_e(x, y), x) - Derivative(x_e(x, y), y)))'), + ('Eq(g, -e.laplace)', 'Eq(e, -g.laplace)', + 'Eq(y_e(x, y), h_x*h_y*x_g(x, y))', + 'Eq(y_g(x, y), h_x*h_y*x_e(x, y))'), + ('Eq(e + g, e.dx + 2.*g.dx)', 'Eq(g + e, g.dx + 2.*e.dx)', + 'Eq(y_e(x, y), h_x*h_y*(x_g(x, y) - 2.0*Derivative(x_g(x, y), x)))', + 'Eq(y_g(x, y), h_x*h_y*(x_e(x, y) - 2.0*Derivative(x_e(x, y), x)))'), + ]) + @skipif('petsc') + def test_coupling(self, eq1, eq2, j01_matvec, j10_matvec): + """ + Test linear coupling between fields, where the off-diagonal + Jacobian submatrices are nonzero. + """ + grid = Grid(shape=(9, 9), dtype=np.float64) - # TODO: FIX THIS LOGIC - THE OUTPUT FOR J01_MATVEC IS Eq(y_e(x, y), -h_x*h_y*x_e(x, y)) BUT IT SHOULD - # BE Eq(y_e(x, y), -h_x*h_y*x_G(x, y)) - - # @pytest.mark.parametrize('eq1, eq2, so, j01_matvec, j10_matvec', [ - # ('Eq(-e.laplace, g)', 'Eq(-g.laplace, e)', '2', 'h_x*h_y*(-2.0/h_y**2 - 2.0/h_x**2)'), - # ]) - # @skipif('petsc') - # def test_coupling(self, eq1, eq2, so, scale): - # """ - # Test linear coupling between fields, where the off-diagonal - # Jacobian submatrices are nonzero. - # """ - # grid = Grid(shape=(9,9), dtype=np.float64) - - # functions = [Function(name=n, grid=grid, space_order=eval(so)) - # for n in ['e', 'f', 'g', 'h']] - # e, f, g, h = functions - - # eq1 = eval(eq1) - # eq2 = eval(eq2) + e = Function(name='e', grid=grid, space_order=2) + g = Function(name='g', grid=grid, space_order=2) - # petsc = PETScSolve({e: [eq1], g: [eq2]}) + eq1 = eval(eq1) + eq2 = eval(eq2) - # jacobian = petsc[0].rhs.fielddata.jacobian + petsc = PETScSolve({e: [eq1], g: [eq2]}) - # j01 = jacobian.get_submatrix(0, 1) - # j10 = jacobian.get_submatrix(1, 0) + jacobian = petsc[0].rhs.fielddata.jacobian - # assert len(j01.matvecs) == 1 - # assert len(j10.matvecs) == 1 + j01 = jacobian.get_submatrix(0, 1) + j10 = jacobian.get_submatrix(1, 0) - # assert j01.col_target == g - # assert j10.col_target == e + assert j01.col_target == g + assert j10.col_target == e - # assert str(j01.matvecs[0]) == j01_matvec - # assert str(j10.matvecs[0]) == j10_matvec + assert str(j01.matvecs[0]) == j01_matvec + assert str(j10.matvecs[0]) == j10_matvec @pytest.mark.parametrize('eq1, eq2, so, scale', [ ('Eq(e.laplace, f)', 'Eq(g.laplace, h)', '2', '-2.0*h_x/h_x**2'), ('Eq(e.laplace, f)', 'Eq(g.laplace, h)', '4', '-2.5*h_x/h_x**2'), ('Eq(e.laplace + e, f)', 'Eq(g.laplace + g, h)', '2', 'h_x*(1 - 2.0/h_x**2)'), ('Eq(e.laplace + e, f)', 'Eq(g.laplace + g, h)', '4', 'h_x*(1 - 2.5/h_x**2)'), - ('Eq(e.laplace + 5.*e, f)', 'Eq(g.laplace + 5.*g, h)', '2', 'h_x*(5.0 - 2.0/h_x**2)'), - ('Eq(e.laplace + 5.*e, f)', 'Eq(g.laplace + 5.*g, h)', '4', 'h_x*(5.0 - 2.5/h_x**2)'), - ('Eq(e.dx + e + e.laplace, f)', 'Eq(g.dx + g + g.laplace, h.dx)', '2', 'h_x*(1 + 1/h_x - 2.0/h_x**2)'), - ('Eq(e.dx + e + e.laplace, f)', 'Eq(g.dx + g + g.laplace, h.dx)', '4', 'h_x*(1 - 2.5/h_x**2)'), - ('Eq(2.*e.laplace + e, f)', 'Eq(2*g.laplace + g, h)', '2', 'h_x*(1 - 4.0/h_x**2)'), - ('Eq(2.*e.laplace + e, f)', 'Eq(2*g.laplace + g, h)', '4', 'h_x*(1 - 5.0/h_x**2)'), + ('Eq(e.laplace + 5.*e, f)', 'Eq(g.laplace + 5.*g, h)', '2', + 'h_x*(5.0 - 2.0/h_x**2)'), + ('Eq(e.laplace + 5.*e, f)', 'Eq(g.laplace + 5.*g, h)', '4', + 'h_x*(5.0 - 2.5/h_x**2)'), + ('Eq(e.dx + e + e.laplace, f)', 'Eq(g.dx + g + g.laplace, h.dx)', '2', + 'h_x*(1 + 1/h_x - 2.0/h_x**2)'), + ('Eq(e.dx + e + e.laplace, f)', 'Eq(g.dx + g + g.laplace, h.dx)', '4', + 'h_x*(1 - 2.5/h_x**2)'), + ('Eq(2.*e.laplace + e, f)', 'Eq(2*g.laplace + g, h)', '2', + 'h_x*(1 - 4.0/h_x**2)'), + ('Eq(2.*e.laplace + e, f)', 'Eq(2*g.laplace + g, h)', '4', + 'h_x*(1 - 5.0/h_x**2)'), ]) @skipif('petsc') def test_jacobian_scaling_1D(self, eq1, eq2, so, scale): """ Test the computation of diagonal scaling in a 1D Jacobian system. - This scaling would be applied to the boundary rows of the matrix - if essential boundary conditions were enforced in the solver. + This scaling would be applied to the boundary rows of the matrix + if essential boundary conditions were enforced in the solver. Its purpose is to reduce the condition number of the matrix. """ grid = Grid(shape=(9,), dtype=np.float64) @@ -931,28 +952,37 @@ def test_jacobian_scaling_1D(self, eq1, eq2, so, scale): assert str(j11.scdiag) == scale @pytest.mark.parametrize('eq1, eq2, so, scale', [ - ('Eq(e.laplace, f)', 'Eq(g.laplace, h)', '2', 'h_x*h_y*(-2.0/h_y**2 - 2.0/h_x**2)'), - ('Eq(e.laplace, f)', 'Eq(g.laplace, h)', '4', 'h_x*h_y*(-2.5/h_y**2 - 2.5/h_x**2)'), - ('Eq(e.laplace + e, f)', 'Eq(g.laplace + g, h)', '2', 'h_x*h_y*(1 - 2.0/h_y**2 - 2.0/h_x**2)'), - ('Eq(e.laplace + e, f)', 'Eq(g.laplace + g, h)', '4', 'h_x*h_y*(1 - 2.5/h_y**2 - 2.5/h_x**2)'), - ('Eq(e.laplace + 5.*e, f)', 'Eq(g.laplace + 5.*g, h)', '2', 'h_x*h_y*(5.0 - 2.0/h_y**2 - 2.0/h_x**2)'), - ('Eq(e.laplace + 5.*e, f)', 'Eq(g.laplace + 5.*g, h)', '4', 'h_x*h_y*(5.0 - 2.5/h_y**2 - 2.5/h_x**2)'), - ('Eq(e.dx + e.dy + e + e.laplace, f)', 'Eq(g.dx + g.dy + g + g.laplace, h)', '2', 'h_x*h_y*(1 + 1/h_y - 2.0/h_y**2 + 1/h_x - 2.0/h_x**2)'), - ('Eq(e.dx + e.dy + e + e.laplace, f)', 'Eq(g.dx + g.dy + g + g.laplace, h)', '4', 'h_x*h_y*(1 - 2.5/h_y**2 - 2.5/h_x**2)'), - ('Eq(2.*e.laplace + e, f)', 'Eq(2*g.laplace + g, h)', '2', 'h_x*h_y*(1 - 4.0/h_y**2 - 4.0/h_x**2)'), - ('Eq(2.*e.laplace + e, f)', 'Eq(2*g.laplace + g, h)', '4', 'h_x*h_y*(1 - 5.0/h_y**2 - 5.0/h_x**2)'), - + ('Eq(e.laplace, f)', 'Eq(g.laplace, h)', '2', + 'h_x*h_y*(-2.0/h_y**2 - 2.0/h_x**2)'), + ('Eq(e.laplace, f)', 'Eq(g.laplace, h)', '4', + 'h_x*h_y*(-2.5/h_y**2 - 2.5/h_x**2)'), + ('Eq(e.laplace + e, f)', 'Eq(g.laplace + g, h)', '2', + 'h_x*h_y*(1 - 2.0/h_y**2 - 2.0/h_x**2)'), + ('Eq(e.laplace + e, f)', 'Eq(g.laplace + g, h)', '4', + 'h_x*h_y*(1 - 2.5/h_y**2 - 2.5/h_x**2)'), + ('Eq(e.laplace + 5.*e, f)', 'Eq(g.laplace + 5.*g, h)', '2', + 'h_x*h_y*(5.0 - 2.0/h_y**2 - 2.0/h_x**2)'), + ('Eq(e.laplace + 5.*e, f)', 'Eq(g.laplace + 5.*g, h)', '4', + 'h_x*h_y*(5.0 - 2.5/h_y**2 - 2.5/h_x**2)'), + ('Eq(e.dx + e.dy + e + e.laplace, f)', 'Eq(g.dx + g.dy + g + g.laplace, h)', + '2', 'h_x*h_y*(1 + 1/h_y - 2.0/h_y**2 + 1/h_x - 2.0/h_x**2)'), + ('Eq(e.dx + e.dy + e + e.laplace, f)', 'Eq(g.dx + g.dy + g + g.laplace, h)', + '4', 'h_x*h_y*(1 - 2.5/h_y**2 - 2.5/h_x**2)'), + ('Eq(2.*e.laplace + e, f)', 'Eq(2*g.laplace + g, h)', '2', + 'h_x*h_y*(1 - 4.0/h_y**2 - 4.0/h_x**2)'), + ('Eq(2.*e.laplace + e, f)', 'Eq(2*g.laplace + g, h)', '4', + 'h_x*h_y*(1 - 5.0/h_y**2 - 5.0/h_x**2)'), ]) @skipif('petsc') def test_jacobian_scaling_2D(self, eq1, eq2, so, scale): """ Test the computation of diagonal scaling in a 2D Jacobian system. - This scaling would be applied to the boundary rows of the matrix - if essential boundary conditions were enforced in the solver. + This scaling would be applied to the boundary rows of the matrix + if essential boundary conditions were enforced in the solver. Its purpose is to reduce the condition number of the matrix. """ - grid = Grid(shape=(9,9), dtype=np.float64) + grid = Grid(shape=(9, 9), dtype=np.float64) functions = [Function(name=n, grid=grid, space_order=eval(so)) for n in ['e', 'f', 'g', 'h']] @@ -972,28 +1002,40 @@ def test_jacobian_scaling_2D(self, eq1, eq2, so, scale): assert str(j11.scdiag) == scale @pytest.mark.parametrize('eq1, eq2, so, scale', [ - ('Eq(e.laplace, f)', 'Eq(g.laplace, h)', '2', 'h_x*h_y*h_z*(-2.0/h_z**2 - 2.0/h_y**2 - 2.0/h_x**2)'), - ('Eq(e.laplace, f)', 'Eq(g.laplace, h)', '4', 'h_x*h_y*h_z*(-2.5/h_z**2 - 2.5/h_y**2 - 2.5/h_x**2)'), - ('Eq(e.laplace + e, f)', 'Eq(g.laplace + g, h)', '2', 'h_x*h_y*h_z*(1 - 2.0/h_z**2 - 2.0/h_y**2 - 2.0/h_x**2)'), - ('Eq(e.laplace + e, f)', 'Eq(g.laplace + g, h)', '4', 'h_x*h_y*h_z*(1 - 2.5/h_z**2 - 2.5/h_y**2 - 2.5/h_x**2)'), - ('Eq(e.laplace + 5.*e, f)', 'Eq(g.laplace + 5.*g, h)', '2', 'h_x*h_y*h_z*(5.0 - 2.0/h_z**2 - 2.0/h_y**2 - 2.0/h_x**2)'), - ('Eq(e.laplace + 5.*e, f)', 'Eq(g.laplace + 5.*g, h)', '4', 'h_x*h_y*h_z*(5.0 - 2.5/h_z**2 - 2.5/h_y**2 - 2.5/h_x**2)'), - ('Eq(e.dx + e.dy + e.dz + e + e.laplace, f)', 'Eq(g.dx + g.dy + g.dz + g + g.laplace, h)', '2', 'h_x*h_y*h_z*(1 + 1/h_z - 2.0/h_z**2 + 1/h_y - 2.0/h_y**2 + 1/h_x - 2.0/h_x**2)'), - ('Eq(e.dx + e.dy + e.dz + e + e.laplace, f)', 'Eq(g.dx + g.dy + g.dz + g + g.laplace, h)', '4', 'h_x*h_y*h_z*(1 - 2.5/h_z**2 - 2.5/h_y**2 - 2.5/h_x**2)'), - ('Eq(2.*e.laplace + e, f)', 'Eq(2*g.laplace + g, h)', '2', 'h_x*h_y*h_z*(1 - 4.0/h_z**2 - 4.0/h_y**2 - 4.0/h_x**2)'), - ('Eq(2.*e.laplace + e, f)', 'Eq(2*g.laplace + g, h)', '4', 'h_x*h_y*h_z*(1 - 5.0/h_z**2 - 5.0/h_y**2 - 5.0/h_x**2)'), - + ('Eq(e.laplace, f)', 'Eq(g.laplace, h)', '2', + 'h_x*h_y*h_z*(-2.0/h_z**2 - 2.0/h_y**2 - 2.0/h_x**2)'), + ('Eq(e.laplace, f)', 'Eq(g.laplace, h)', '4', + 'h_x*h_y*h_z*(-2.5/h_z**2 - 2.5/h_y**2 - 2.5/h_x**2)'), + ('Eq(e.laplace + e, f)', 'Eq(g.laplace + g, h)', '2', + 'h_x*h_y*h_z*(1 - 2.0/h_z**2 - 2.0/h_y**2 - 2.0/h_x**2)'), + ('Eq(e.laplace + e, f)', 'Eq(g.laplace + g, h)', '4', + 'h_x*h_y*h_z*(1 - 2.5/h_z**2 - 2.5/h_y**2 - 2.5/h_x**2)'), + ('Eq(e.laplace + 5.*e, f)', 'Eq(g.laplace + 5.*g, h)', '2', + 'h_x*h_y*h_z*(5.0 - 2.0/h_z**2 - 2.0/h_y**2 - 2.0/h_x**2)'), + ('Eq(e.laplace + 5.*e, f)', 'Eq(g.laplace + 5.*g, h)', '4', + 'h_x*h_y*h_z*(5.0 - 2.5/h_z**2 - 2.5/h_y**2 - 2.5/h_x**2)'), + ('Eq(e.dx + e.dy + e.dz + e + e.laplace, f)', + 'Eq(g.dx + g.dy + g.dz + g + g.laplace, h)', '2', + 'h_x*h_y*h_z*(1 + 1/h_z - 2.0/h_z**2 + 1/h_y - 2.0/h_y**2 + ' + + '1/h_x - 2.0/h_x**2)'), + ('Eq(e.dx + e.dy + e.dz + e + e.laplace, f)', + 'Eq(g.dx + g.dy + g.dz + g + g.laplace, h)', '4', + 'h_x*h_y*h_z*(1 - 2.5/h_z**2 - 2.5/h_y**2 - 2.5/h_x**2)'), + ('Eq(2.*e.laplace + e, f)', 'Eq(2*g.laplace + g, h)', '2', + 'h_x*h_y*h_z*(1 - 4.0/h_z**2 - 4.0/h_y**2 - 4.0/h_x**2)'), + ('Eq(2.*e.laplace + e, f)', 'Eq(2*g.laplace + g, h)', '4', + 'h_x*h_y*h_z*(1 - 5.0/h_z**2 - 5.0/h_y**2 - 5.0/h_x**2)'), ]) @skipif('petsc') def test_jacobian_scaling_3D(self, eq1, eq2, so, scale): """ Test the computation of diagonal scaling in a 3D Jacobian system. - This scaling would be applied to the boundary rows of the matrix - if essential boundary conditions were enforced in the solver. + This scaling would be applied to the boundary rows of the matrix + if essential boundary conditions were enforced in the solver. Its purpose is to reduce the condition number of the matrix. """ - grid = Grid(shape=(9,9,9), dtype=np.float64) + grid = Grid(shape=(9, 9, 9), dtype=np.float64) functions = [Function(name=n, grid=grid, space_order=eval(so)) for n in ['e', 'f', 'g', 'h']] @@ -1012,9 +1054,6 @@ def test_jacobian_scaling_3D(self, eq1, eq2, so, scale): assert str(j00.scdiag) == scale assert str(j11.scdiag) == scale - # test coupled residual callback - check the .dot for each field etc - - @skipif('petsc') def test_residual_bundle(self): grid = Grid(shape=(11, 11), dtype=np.float64) @@ -1055,6 +1094,42 @@ def test_residual_bundle(self): assert 'struct Field0\n{\n PetscScalar e;\n PetscScalar f;\n ' \ + 'PetscScalar g;\n}' in str(op3.ccode) + @skipif('petsc') + def test_residual_callback(self): + """ + Check that the main residual callback correctly accesses the + target fields in the bundle. + """ + grid = Grid(shape=(9, 9), dtype=np.float64) + + functions = [Function(name=n, grid=grid, space_order=2) + for n in ['e', 'f', 'g', 'h']] + e, f, g, h = functions + + eq1 = Eq(e.laplace, f) + eq2 = Eq(g.laplace, h) + + petsc = PETScSolve({e: [eq1], g: [eq2]}) + + with switchconfig(language='petsc'): + op = Operator(petsc) + + # Check the residual callback + residual = op._func_table['WholeFormFunc0'].root + + exprs = FindNodes(Expression).visit(residual) + exprs = [str(e) for e in exprs] + + assert 'f_bundle[x + 2][y + 2].e = (r4*x_bundle[x + 1][y + 2].e + ' + \ + 'r4*x_bundle[x + 3][y + 2].e + r5*x_bundle[x + 2][y + 1].e + r5*' + \ + 'x_bundle[x + 2][y + 3].e - 2.0*(r4*x_bundle[x + 2][y + 2].e + r5*' + \ + 'x_bundle[x + 2][y + 2].e) - f[x + 2][y + 2])*ctx0->h_x*ctx0->h_y;' in exprs + + assert 'f_bundle[x + 2][y + 2].g = (r4*x_bundle[x + 1][y + 2].g + ' + \ + 'r4*x_bundle[x + 3][y + 2].g + r5*x_bundle[x + 2][y + 1].g + r5*' + \ + 'x_bundle[x + 2][y + 3].g - 2.0*(r4*x_bundle[x + 2][y + 2].g + r5*' + \ + 'x_bundle[x + 2][y + 2].g) - h[x + 2][y + 2])*ctx0->h_x*ctx0->h_y;' in exprs + @skipif('petsc') def test_essential_bcs(self): """ From b6710be878ddbf51d9d225af4bba31b6284d25bb Mon Sep 17 00:00:00 2001 From: ZoeLeibowitz Date: Wed, 11 Jun 2025 00:03:07 +0100 Subject: [PATCH 42/84] types: Fix PetscBundle symbolic shape --- devito/petsc/types/array.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/devito/petsc/types/array.py b/devito/petsc/types/array.py index c7a09f8265..e16b036c5a 100644 --- a/devito/petsc/types/array.py +++ b/devito/petsc/types/array.py @@ -140,6 +140,10 @@ def __init__(self, *args, pname="Field", **kwargs): def _C_ctype(self): fields = [(i.target.name, dtype_to_ctype(i.dtype)) for i in self.components] return POINTER(type(self.pname, (Structure,), {'_fields_': fields})) + + @cached_property + def symbolic_shape(self): + return self.c0.symbolic_shape @cached_property def indexed(self): From 01d153ad8f3fb57eb97afcffa04f03456312c47a Mon Sep 17 00:00:00 2001 From: ZoeLeibowitz Date: Wed, 11 Jun 2025 00:50:32 +0100 Subject: [PATCH 43/84] tests: Fix petsc tests --- tests/test_petsc.py | 12 ++++++------ 1 file changed, 6 insertions(+), 6 deletions(-) diff --git a/tests/test_petsc.py b/tests/test_petsc.py index 7fd59404d3..d5087c13e6 100644 --- a/tests/test_petsc.py +++ b/tests/test_petsc.py @@ -621,12 +621,12 @@ class TestCoupledLinear: ('Eq(e.laplace, f)', 'Eq(g.laplace, h)', '2'), ('Eq(e.laplace, f)', 'Eq(g.laplace, h)', '4'), ('Eq(e.laplace, f)', 'Eq(g.laplace, h)', '6'), - ('Eq(e.laplace, f + 5)', 'Eq(g.laplace, h + 5)', '2'), - ('Eq(e.laplace, f + 5)', 'Eq(g.laplace, h + 5)', '4'), - ('Eq(e.laplace, f + 5)', 'Eq(g.laplace, h + 5)', '6'), - ('Eq(e.dx, e + 2*f)', 'Eq(g.dx, g + 2*h)', '2'), - ('Eq(e.dx, e + 2*f)', 'Eq(g.dx, g + 2*h)', '4'), - ('Eq(e.dx, e + 2*f)', 'Eq(g.dx, g + 2*h)', '6'), + ('Eq(e.laplace, f + 5.)', 'Eq(g.laplace, h + 5.)', '2'), + ('Eq(e.laplace, f + 5.)', 'Eq(g.laplace, h + 5.)', '4'), + ('Eq(e.laplace, f + 5.)', 'Eq(g.laplace, h + 5.)', '6'), + ('Eq(e.dx, e + 2.*f)', 'Eq(g.dx, g + 2.*h)', '2'), + ('Eq(e.dx, e + 2.*f)', 'Eq(g.dx, g + 2.*h)', '4'), + ('Eq(e.dx, e + 2.*f)', 'Eq(g.dx, g + 2.*h)', '6'), ('Eq(f.dx, e.dx + e + e.laplace)', 'Eq(h.dx, g.dx + g + g.laplace)', '2'), ('Eq(f.dx, e.dx + e + e.laplace)', 'Eq(h.dx, g.dx + g + g.laplace)', '4'), ('Eq(f.dx, e.dx + e + e.laplace)', 'Eq(h.dx, g.dx + g + g.laplace)', '6'), From d6bd4f4c96362ce4398c1e8c7606877f152057dd Mon Sep 17 00:00:00 2001 From: ZoeLeibowitz Date: Wed, 11 Jun 2025 11:57:42 +0100 Subject: [PATCH 44/84] misc: Address comments --- devito/passes/iet/engine.py | 1 - devito/petsc/iet/routines.py | 52 +++++++++++++++++++--------------- devito/petsc/solve.py | 2 +- devito/petsc/types/array.py | 40 ++++++++++++++++++++++---- devito/petsc/types/types.py | 51 ++++++++++++++++++--------------- devito/symbolics/extraction.py | 2 ++ devito/types/array.py | 6 ++-- tests/test_petsc.py | 2 +- 8 files changed, 100 insertions(+), 56 deletions(-) diff --git a/devito/passes/iet/engine.py b/devito/passes/iet/engine.py index ad4377d979..0dd6ad0a56 100644 --- a/devito/passes/iet/engine.py +++ b/devito/passes/iet/engine.py @@ -530,7 +530,6 @@ def _(i, mapper, sregistry): }) -@abstract_object.register(Array) @abstract_object.register(ArrayBasic) def _(i, mapper, sregistry): if isinstance(i, Lock): diff --git a/devito/petsc/iet/routines.py b/devito/petsc/iet/routines.py index 4e48a571f2..3ddf78e770 100644 --- a/devito/petsc/iet/routines.py +++ b/devito/petsc/iet/routines.py @@ -114,6 +114,10 @@ def fielddata(self): def arrays(self): return self.fielddata.arrays + @property + def target(self): + return self.fielddata.target + def _make_core(self): self._make_matvec(self.fielddata.jacobian) self._make_formfunc() @@ -163,9 +167,7 @@ def _create_matvec_body(self, body, jacobian): 'DMGetApplicationContext', [dmda, Byref(ctx._C_symbol)] ) - zero_y_memory = petsc_call( - 'VecSet', [objs['Y'], 0.0] - ) if self.zero_memory else None + zero_y_memory = self.zero_vector(objs['Y']) dm_get_local_xvec = petsc_call( 'DMGetLocalVector', [dmda, Byref(xlocal)] @@ -297,7 +299,8 @@ def _create_formfunc_body(self, body): linsolve_expr = self.injectsolve.expr.rhs objs = self.objs sobjs = self.solver_objs - target = self.fielddata.target + arrays = self.arrays + target = self.target dmda = sobjs['callbackdm'] ctx = objs['dummyctx'] @@ -307,8 +310,8 @@ def _create_formfunc_body(self, body): fields = self._dummy_fields(body) self._struct_params.extend(fields) - f_formfunc = self.fielddata.arrays[target]['f'] - x_formfunc = self.fielddata.arrays[target]['x'] + f_formfunc = arrays[target]['f'] + x_formfunc = arrays[target]['x'] dm_cast = DummyExpr(dmda, DMCast(objs['dummyptr']), init=True) @@ -316,9 +319,7 @@ def _create_formfunc_body(self, body): 'DMGetApplicationContext', [dmda, Byref(ctx._C_symbol)] ) - zero_f_memory = petsc_call( - 'VecSet', [objs['F'], 0.0] - ) if self.zero_memory else None + zero_f_memory = self.zero_vector(objs['F']) dm_get_local_xvec = petsc_call( 'DMGetLocalVector', [dmda, Byref(objs['xloc'])] @@ -437,7 +438,7 @@ def _create_form_rhs_body(self, body): linsolve_expr = self.injectsolve.expr.rhs objs = self.objs sobjs = self.solver_objs - target = self.fielddata.target + target = self.target dmda = sobjs['callbackdm'] ctx = objs['dummyctx'] @@ -508,13 +509,15 @@ def _create_form_rhs_body(self, body): ) # Dereference function data in struct - dereference_funcs = [Dereference(i, ctx) for i in - fields if isinstance(i.function, AbstractFunction)] + dereference_funcs = tuple( + [Dereference(i, ctx) for i in + fields if isinstance(i.function, AbstractFunction)] + ) formrhs_body = CallableBody( List(body=[body]), init=(objs['begin_user'],), - stacks=stacks+tuple(dereference_funcs), + stacks=stacks+dereference_funcs, retstmt=(Call('PetscFunctionReturn', arguments=[0]),) ) @@ -550,7 +553,7 @@ def _create_initial_guess_body(self, body): linsolve_expr = self.injectsolve.expr.rhs objs = self.objs sobjs = self.solver_objs - target = self.fielddata.target + target = self.target dmda = sobjs['callbackdm'] ctx = objs['dummyctx'] @@ -653,6 +656,12 @@ def _uxreplace_efuncs(self): mapper.update({k: visitor.visit(v)}) return mapper + def zero_vector(self, vec): + """ + Zeros the memory of the output vector before computation + """ + return petsc_call('VecSet', [vec, 0.0]) if self.zero_memory else None + class CCBBuilder(CBBuilder): def __init__(self, **kwargs): @@ -1059,7 +1068,7 @@ def residual_bundle(self, body, bundles): if i.base in mapper: bundle = mapper[i.base] index = bundles['target_indices'][i.function.target] - index = (index,)+i.indices + index = (index,) + i.indices subs[i] = bundle.__getitem__(index) body = Uxreplace(subs).visit(body) @@ -1178,9 +1187,9 @@ def _extend_build(self, base_dict): base_dict[f'{name}F'] = CallbackVec(f'{name}F') # Bundle objects/metadata required by the coupled residual callback - f_components = [] - x_components = [] + f_components, x_components = [], [] bundle_mapper = {} + pname = sreg.make_name(prefix='Field') target_indices = {t: i for i, t in enumerate(targets)} @@ -1190,18 +1199,15 @@ def _extend_build(self, base_dict): f_components.append(f_arr) x_components.append(x_arr) - bundle_pname = sreg.make_name(prefix='Field') fbundle = PetscBundle( - name='f_bundle', components=f_components, pname=bundle_pname + name='f_bundle', components=f_components, pname=pname ) xbundle = PetscBundle( - name='x_bundle', components=x_components, pname=bundle_pname + name='x_bundle', components=x_components, pname=pname ) # Build the bundle mapper - for i, t in enumerate(targets): - f_arr = arrays[t]['f'] - x_arr = arrays[t]['x'] + for f_arr, x_arr in zip(f_components, x_components): bundle_mapper[f_arr.base] = fbundle bundle_mapper[x_arr.base] = xbundle diff --git a/devito/petsc/solve.py b/devito/petsc/solve.py index 61d0c187e6..9a85774ef8 100644 --- a/devito/petsc/solve.py +++ b/devito/petsc/solve.py @@ -83,7 +83,7 @@ def linear_solve_args(self): self.time_mapper = generate_time_mapper(funcs) arrays = self.generate_arrays_combined(target) - eqns = sorted(eqns, key=lambda e: 0 if isinstance(e, EssentialBC) else 1) + eqns = sorted(eqns, key=lambda e: not isinstance(e, EssentialBC)) jacobian = Jacobian(target, eqns, arrays, self.time_mapper) residual = Residual(target, eqns, arrays, self.time_mapper, jacobian.scdiag) diff --git a/devito/petsc/types/array.py b/devito/petsc/types/array.py index e16b036c5a..e691a96c33 100644 --- a/devito/petsc/types/array.py +++ b/devito/petsc/types/array.py @@ -126,6 +126,34 @@ def symbolic_shape(self): class PetscBundle(Bundle): """ + Tensor symbol representing an unrolled vector of PETScArrays. + + This class declares a struct in the generated ccode to represent the + fields defined at each node of the grid. For example: + + typedef struct { + PetscScalar u,v,omega,temperature; + } Field; + + Residual evaluations are then written using: + + f[i][j].omega = ... + + Reference - https://petsc.org/release/manual/vec/#sec-struct + + Parameters + ---------- + name : str + Name of the symbol. + components : tuple of PETScArray + The PETScArrays of the Bundle. + pname : str, optional + The name of the struct in the generated C code. Defaults to "Field". + + Warnings + -------- + PetscBundles are created and managed directly by Devito (IOW, they are not + expected to be used directly in user code). """ is_Bundle = True _data_alignment = False @@ -136,11 +164,11 @@ def __init__(self, *args, pname="Field", **kwargs): super().__init__(*args, **kwargs) self._pname = pname - @property + @cached_property def _C_ctype(self): fields = [(i.target.name, dtype_to_ctype(i.dtype)) for i in self.components] return POINTER(type(self.pname, (Structure,), {'_fields_': fields})) - + @cached_property def symbolic_shape(self): return self.c0.symbolic_shape @@ -176,8 +204,10 @@ def __getitem__(self, index): component_names=names ) else: - raise ValueError("Expected %d or %d indices, got %d instead" - % (self.ndim, self.ndim + 1, len(index))) + raise ValueError( + f"Expected {self.ndim} or {self.ndim + 1} indices, " + f"got {len(index)} instead" + ) @property def pname(self): @@ -187,7 +217,7 @@ def pname(self): class PetscComponentAccess(ComponentAccess): def __new__(cls, arg, index=0, component_names=None, **kwargs): if not arg.is_Indexed: - raise ValueError("Expected Indexed, got `%s` instead" % type(arg)) + raise ValueError(f"Expected Indexed, got `{type(arg)}` instead") names = component_names or cls._default_component_names obj = Expr.__new__(cls, arg) diff --git a/devito/petsc/types/types.py b/devito/petsc/types/types.py index 44d43d2f97..2e866896cb 100644 --- a/devito/petsc/types/types.py +++ b/devito/petsc/types/types.py @@ -2,7 +2,7 @@ from itertools import chain -from devito.tools import Reconstructable, sympy_mutex, as_tuple +from devito.tools import Reconstructable, sympy_mutex, as_tuple, frozendict from devito.tools.dtypes_lowering import dtype_mapper from devito.petsc.utils import petsc_variables from devito.symbolics.extraction import separate_eqn, generate_targets, centre_stencil @@ -269,11 +269,11 @@ def col_target(self): return self.target def _build_matvecs(self): - matvecs = [ - e for eq in self.eqns for e in - self._build_matvec_eq(eq) - if e is not None - ] + matvecs = [] + for eq in self.eqns: + matvecs.extend( + e for e in self._build_matvec_eq(eq) if e is not None + ) matvecs = tuple(sorted(matvecs, key=lambda e: not isinstance(e, EssentialBC))) matvecs = self._scale_non_bcs(matvecs) @@ -287,12 +287,13 @@ def _build_matvec_eq(self, eq, col_target=None, row_target=None): col_target = col_target or self.target row_target = row_target or self.target - b, F_target, _, targets = separate_eqn(eq, col_target) + _, F_target, _, targets = separate_eqn(eq, col_target) if F_target: return self._make_matvec( eq, F_target, targets, col_target, row_target ) - return (None,) + else: + return (None,) def _make_matvec(self, eq, F_target, targets, col_target, row_target): y = self.arrays[row_target]['y'] @@ -307,7 +308,7 @@ def _make_matvec(self, eq, F_target, targets, col_target, row_target): else: rhs = F_target.subs(targets_to_arrays(x, targets)) rhs = rhs.subs(self.time_mapper) - return as_tuple(Eq(y, rhs, subdomain=eq.subdomain)) + return (Eq(y, rhs, subdomain=eq.subdomain),) def _scale_non_bcs(self, matvecs, target=None): target = target or self.target @@ -376,7 +377,7 @@ class MixedJacobian(Jacobian): def __init__(self, target_eqns, arrays, time_mapper): """ """ - self.targets = as_tuple(target_eqns.keys()) + self.targets = tuple(target_eqns.keys()) self.arrays = arrays self.time_mapper = time_mapper self._submatrices = [] @@ -391,7 +392,7 @@ def submatrices(self): return self._submatrices @property - def no_submatrices(self): + def n_submatrices(self): """ Return the number of submatrix blocks. """ @@ -421,11 +422,13 @@ def _build_blocks(self, target_eqns): for i, row_target in enumerate(self.targets): eqns = target_eqns[row_target] for j, col_target in enumerate(self.targets): - matvecs = [ - e for eq in eqns for e in - self._build_matvec_eq(eq, col_target, row_target) - ] + matvecs = [] + for eq in eqns: + matvecs.extend( + e for e in self._build_matvec_eq(eq, col_target, row_target) + ) matvecs = [m for m in matvecs if m is not None] + # Sort to put EssentialBC first if any matvecs = tuple( sorted(matvecs, key=lambda e: not isinstance(e, EssentialBC)) @@ -461,7 +464,7 @@ def __repr__(self): f"{sm.name} (row={sm.row_idx}, col={sm.col_idx})" for sm in self.submatrices ) - return f"" + return f"" class Residual: @@ -499,12 +502,13 @@ def _build_equations(self): # TODO: If b is zero then don't need a rhs vector+callback rhs.extend(self._make_b(eq, b)) - self._formfuncs = [self._scale_bcs(eq) for eq in funcs] - self._formrhs = rhs + self._formfuncs = tuple([self._scale_bcs(eq) for eq in funcs]) + self._formrhs = tuple(rhs) def _make_F_target(self, eq, F_target, targets): arrays = self.arrays[self.target] volume = self.target.grid.symbolic_volume_cell + if isinstance(eq, EssentialBC): # The initial guess satisfies the essential BCs, so this term is zero. # Still included to support Jacobian testing via finite differences. @@ -513,19 +517,20 @@ def _make_F_target(self, eq, F_target, targets): # Move essential boundary condition to the right-hand side zero_col = ZeroColumn(arrays['x'], eq.rhs, subdomain=eq.subdomain) return (zero_row, zero_col) + else: if isinstance(F_target, (int, float)): rhs = F_target * volume else: rhs = F_target.subs(targets_to_arrays(arrays['x'], targets)) rhs = rhs.subs(self.time_mapper) * volume - return as_tuple(Eq(arrays['f'], rhs, subdomain=eq.subdomain)) + return (Eq(arrays['f'], rhs, subdomain=eq.subdomain),) def _make_b(self, eq, b): b_arr = self.arrays[self.target]['b'] rhs = 0. if isinstance(eq, EssentialBC) else b.subs(self.time_mapper) rhs = rhs * self.target.grid.symbolic_volume_cell - return as_tuple(Eq(b_arr, rhs, subdomain=eq.subdomain)) + return (Eq(b_arr, rhs, subdomain=eq.subdomain),) def _scale_bcs(self, eq, scdiag=None): """ @@ -617,11 +622,11 @@ def _build_equations(self): """ Return a list of initial guess equations. """ - self._eqs = [ + self._eqs = tuple([ eq for eq in (self._make_initial_guess(e) for e in self.eqns) if eq is not None - ] + ]) def _make_initial_guess(self, eq): if isinstance(eq, EssentialBC): @@ -655,4 +660,4 @@ def targets_to_arrays(array, targets): array_targets = [ array.subs(dict(zip(array.indices, i))) for i in space_indices ] - return dict(zip(targets, array_targets)) + return frozendict(zip(targets, array_targets)) diff --git a/devito/symbolics/extraction.py b/devito/symbolics/extraction.py index dfbc145748..fc47e94444 100644 --- a/devito/symbolics/extraction.py +++ b/devito/symbolics/extraction.py @@ -17,9 +17,11 @@ def separate_eqn(eqn, target): where F(target) = b. """ zeroed_eqn = Eq(eqn.lhs - eqn.rhs, 0) + from devito.operations.solve import eval_time_derivatives zeroed_eqn = eval_time_derivatives(zeroed_eqn.lhs) target_funcs = set(generate_targets(zeroed_eqn, target)) + b, F_target = remove_targets(zeroed_eqn, target_funcs) return -b, F_target, zeroed_eqn, target_funcs diff --git a/devito/types/array.py b/devito/types/array.py index 8c7fdb8a3d..aec3bd0626 100644 --- a/devito/types/array.py +++ b/devito/types/array.py @@ -531,8 +531,10 @@ def __getitem__(self, index): component_index, indices = index[0], index[1:] return ComponentAccess(self.indexed[indices], component_index) else: - raise ValueError("Expected %d or %d indices, got %d instead" - % (self.ndim, self.ndim + 1, len(index))) + raise ValueError( + f"Expected {self.ndim} or {self.ndim + 1} indices, " + f"got {len(index)} instead" + ) @property def _C_ctype(self): diff --git a/tests/test_petsc.py b/tests/test_petsc.py index d5087c13e6..10e030b064 100644 --- a/tests/test_petsc.py +++ b/tests/test_petsc.py @@ -826,7 +826,7 @@ def test_mixed_jacobian(self): j11 = jacobian.get_submatrix(1, 1) # Check the number of submatrices - assert jacobian.no_submatrices == 4 + assert jacobian.n_submatrices == 4 # Technically a non-coupled problem, so the only non-zero submatrices # should be the diagonal ones i.e J00 and J11 From 4fad6aa52235df66975eb85902b77988aad916b7 Mon Sep 17 00:00:00 2001 From: ZoeLeibowitz Date: Wed, 11 Jun 2025 12:04:04 +0100 Subject: [PATCH 45/84] misc: Merge leftover --- tests/test_symbolics.py | 5 ++--- 1 file changed, 2 insertions(+), 3 deletions(-) diff --git a/tests/test_symbolics.py b/tests/test_symbolics.py index ec90facd13..61c45527f8 100644 --- a/tests/test_symbolics.py +++ b/tests/test_symbolics.py @@ -957,7 +957,6 @@ def test_print_div(): assert cstr == 'sizeof(int)/sizeof(long)' -<<<<<<< HEAD @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))')), @@ -1194,7 +1193,8 @@ def test_centre_stencil(expr, so, target, expected): centre = centre_stencil(eval(expr), eval(target)) assert str(centre) == expected -======= + + def test_customdtype_complex(): """ Test that `CustomDtype` doesn't brak is_imag @@ -1327,4 +1327,3 @@ def test_conj(self): Operator([Eq(g, Conj(f))])() assert np.all(np.isclose(g.data, np.conj(f.data))) ->>>>>>> petsc From 27dd201b0365e3b07895d2dbc47ae7826692ae3e Mon Sep 17 00:00:00 2001 From: ZoeLeibowitz Date: Wed, 11 Jun 2025 17:33:19 +0100 Subject: [PATCH 46/84] compiler: Edit switch_log_level --- devito/logger.py | 5 +++-- devito/operator/operator.py | 6 +++++- devito/petsc/iet/passes.py | 18 +++++++++--------- devito/petsc/iet/routines.py | 14 +++++++++----- devito/petsc/initialize.py | 7 ++++--- 5 files changed, 30 insertions(+), 20 deletions(-) diff --git a/devito/logger.py b/devito/logger.py index 9c1dfa56d2..c9834dbfe7 100644 --- a/devito/logger.py +++ b/devito/logger.py @@ -76,9 +76,10 @@ def set_log_level(level, comm=None): used, for example, if one wants to log to one file per rank. """ from devito import configuration - + from devito.mpi.distributed import MPI + if comm is not None and configuration['mpi']: - if comm.rank != 0: + if comm!= MPI.COMM_NULL and comm.rank != 0: logger.removeHandler(stream_handler) logger.addHandler(logging.NullHandler()) else: diff --git a/devito/operator/operator.py b/devito/operator/operator.py index 3a36952a58..5216bde481 100644 --- a/devito/operator/operator.py +++ b/devito/operator/operator.py @@ -38,7 +38,7 @@ from devito.petsc.iet.passes import lower_petsc from devito.petsc.clusters import petsc_preprocess -__all__ = ['Operator'] +__all__ = ['Operator', 'SpecialOp'] class Operator(Callable): @@ -1456,3 +1456,7 @@ def parse_kwargs(**kwargs): kwargs['subs'] = {k: sympify(v) for k, v in kwargs.get('subs', {}).items()} return kwargs + + +class SpecialOp(Operator): + pass diff --git a/devito/petsc/iet/passes.py b/devito/petsc/iet/passes.py index 3c73bfd391..292f455acf 100644 --- a/devito/petsc/iet/passes.py +++ b/devito/petsc/iet/passes.py @@ -52,9 +52,10 @@ def lower_petsc(iet, **kwargs): # Assumption is that all solves are on the same grid if len(unique_grids) > 1: raise ValueError("All PETScSolves must use the same Grid, but multiple found.") + grid = unique_grids.pop() # Create core PETSc calls (not specific to each PETScSolve) - core = make_core_petsc_calls(objs, **kwargs) + core = make_core_petsc_calls(objs, grid, **kwargs) setup = [] subs = {} @@ -62,7 +63,7 @@ def lower_petsc(iet, **kwargs): for iters, (injectsolve,) in injectsolve_mapper.items(): - builder = Builder(injectsolve, objs, iters, **kwargs) + builder = Builder(injectsolve, objs, iters, grid, **kwargs) setup.extend(builder.solversetup.calls) @@ -108,8 +109,9 @@ def finalize(iet): return iet._rebuild(body=finalize_body) -def make_core_petsc_calls(objs, **kwargs): - call_mpi = petsc_call_mpi('MPI_Comm_size', [objs['comm'], Byref(objs['size'])]) +def make_core_petsc_calls(objs, grid, **kwargs): + comm = grid.distributor._obj_comm + call_mpi = petsc_call_mpi('MPI_Comm_size', [comm, Byref(objs['size'])]) return call_mpi, BlankLine @@ -123,16 +125,18 @@ class Builder: returning subclasses of the objects initialised in __init__, depending on the properties of `injectsolve`. """ - def __init__(self, injectsolve, objs, iters, **kwargs): + def __init__(self, injectsolve, objs, iters, grid, **kwargs): self.injectsolve = injectsolve self.objs = objs self.iters = iters + self.grid = grid self.kwargs = kwargs self.coupled = isinstance(injectsolve.expr.rhs.fielddata, MultipleFieldData) self.args = { 'injectsolve': self.injectsolve, 'objs': self.objs, 'iters': self.iters, + 'grid': self.grid, **self.kwargs } self.args['solver_objs'] = self.objbuilder.solver_objs @@ -190,9 +194,6 @@ def populate_matrix_context(efuncs, objs): ) -# TODO: Devito MPI + PETSc testing -# if kwargs['options']['mpi'] -> communicator = grid.distributor._obj_comm -communicator = 'PETSC_COMM_WORLD' subdms = PointerDM(name='subdms') fields = PointerIS(name='fields') submats = PointerMat(name='submats') @@ -208,7 +209,6 @@ def populate_matrix_context(efuncs, objs): # they are semantically identical. objs = frozendict({ 'size': PetscMPIInt(name='size'), - 'comm': communicator, 'err': PetscErrorCode(name='err'), 'block': CallbackMat('block'), 'submat_arr': PointerMat(name='submat_arr'), diff --git a/devito/petsc/iet/routines.py b/devito/petsc/iet/routines.py index 3ddf78e770..ae305a5ba0 100644 --- a/devito/petsc/iet/routines.py +++ b/devito/petsc/iet/routines.py @@ -963,7 +963,7 @@ def _submat_callback_body(self): ptr = DummyExpr(objs['submat_arr']._C_symbol, Deref(objs['Submats']), init=True) - mat_create = petsc_call('MatCreate', [self.objs['comm'], Byref(objs['block'])]) + mat_create = petsc_call('MatCreate', [sobjs['comm'], Byref(objs['block'])]) mat_set_sizes = petsc_call( 'MatSetSizes', [ @@ -1085,6 +1085,7 @@ def __init__(self, **kwargs): self.injectsolve = kwargs.get('injectsolve') self.objs = kwargs.get('objs') self.sregistry = kwargs.get('sregistry') + self.grid = kwargs.get('grid') self.fielddata = self.injectsolve.expr.rhs.fielddata self.solver_objs = self._build() @@ -1124,6 +1125,8 @@ def _build(self): 'dmda': DM(sreg.make_name(prefix='da'), dofs=len(targets)), 'callbackdm': CallbackDM(sreg.make_name(prefix='dm')), } + # TODO: Devito MPI + PETSc testing + base_dict['comm'] = self.grid.distributor._obj_comm self._target_dependent(base_dict) return self._extend_build(base_dict) @@ -1279,7 +1282,7 @@ def _setup(self): solver_params = self.injectsolve.expr.rhs.solver_parameters - snes_create = petsc_call('SNESCreate', [objs['comm'], Byref(sobjs['snes'])]) + snes_create = petsc_call('SNESCreate', [sobjs['comm'], Byref(sobjs['snes'])]) snes_set_dm = petsc_call('SNESSetDM', [sobjs['snes'], dmda]) @@ -1305,7 +1308,7 @@ def _setup(self): v for v, dim in zip(target.shape_allocated, target.dimensions) if dim.is_Space ) local_x = petsc_call('VecCreateMPIWithArray', - ['PETSC_COMM_WORLD', 1, local_size, 'PETSC_DECIDE', + [sobjs['comm'], 1, local_size, 'PETSC_DECIDE', field_from_ptr, Byref(sobjs['xlocal'])]) # TODO: potentially also need to set the DM and local/global map to xlocal @@ -1408,11 +1411,12 @@ def _create_dmda_calls(self, dmda): def _create_dmda(self, dmda): objs = self.objs + sobjs = self.solver_objs grid = self.fielddata.grid nspace_dims = len(grid.dimensions) # MPI communicator - args = [objs['comm']] + args = [sobjs['comm']] # Type of ghost nodes args.extend(['DM_BOUNDARY_GHOSTED' for _ in range(nspace_dims)]) @@ -1453,7 +1457,7 @@ def _setup(self): solver_params = self.injectsolve.expr.rhs.solver_parameters - snes_create = petsc_call('SNESCreate', [objs['comm'], Byref(sobjs['snes'])]) + snes_create = petsc_call('SNESCreate', [sobjs['comm'], Byref(sobjs['snes'])]) snes_set_dm = petsc_call('SNESSetDM', [sobjs['snes'], dmda]) diff --git a/devito/petsc/initialize.py b/devito/petsc/initialize.py index 80e3c7520c..3fc419c26a 100644 --- a/devito/petsc/initialize.py +++ b/devito/petsc/initialize.py @@ -3,7 +3,8 @@ from ctypes import POINTER, cast, c_char import atexit -from devito import Operator, switchconfig +from devito import switchconfig +from devito.operator.operator import SpecialOp from devito.types import Symbol from devito.types.equation import PetscEq from devito.petsc.types import Initialize, Finalize @@ -21,11 +22,11 @@ def PetscInitialize(): # This would prevent circular imports when initializing during import # from the PETSc module. with switchconfig(language='petsc'): - op_init = Operator( + op_init = SpecialOp( [PetscEq(dummy, Initialize(dummy))], name='kernel_init', opt='noop' ) - op_finalize = Operator( + op_finalize = SpecialOp( [PetscEq(dummy, Finalize(dummy))], name='kernel_finalize', opt='noop' ) From 30ed5bac8e551b00be9c8485f0bf81702d3e7476 Mon Sep 17 00:00:00 2001 From: ZoeLeibowitz Date: Wed, 11 Jun 2025 18:16:45 +0100 Subject: [PATCH 47/84] tests: Add mpi petsc test --- devito/petsc/iet/passes.py | 32 ++++++++++++++++++++++++++++++-- devito/petsc/iet/routines.py | 8 +++++++- examples/petsc/petsc_test.py | 2 +- 3 files changed, 38 insertions(+), 4 deletions(-) diff --git a/devito/petsc/iet/passes.py b/devito/petsc/iet/passes.py index 292f455acf..f77879cf10 100644 --- a/devito/petsc/iet/passes.py +++ b/devito/petsc/iet/passes.py @@ -93,6 +93,17 @@ def initialize(iet): help_string = c.Line(r'static char help[] = "This is help text.\n";') init_body = petsc_call('PetscInitialize', [Byref(argc), Byref(argv), Null, Help]) + + # print_comm_size = + + # size = c.Line('PetscMPIInt size;') + # get_size = c.Line('PetscCallMPI(MPI_Comm_size(comm,&(size)));') + # rank = c.Line('PetscMPIInt rank;') + # get_rank = c.Line('PetscCallMPI(MPI_Comm_rank(comm,&(rank)));') + # print_comm_size = c.Line('PetscSynchronizedPrintf(comm, "MPI_Comm_size: %d\\n", size);') + # print_comm_rank = c.Line('PetscSynchronizedPrintf(comm, "MPI_Comm_rank: %d\\n", rank);') + + init_body = CallableBody( body=(petsc_func_begin_user, help_string, init_body), retstmt=(Call('PetscFunctionReturn', arguments=[0]),) @@ -110,10 +121,27 @@ def finalize(iet): def make_core_petsc_calls(objs, grid, **kwargs): - comm = grid.distributor._obj_comm + devito_mpi = kwargs['options'].get('mpi', False) + if devito_mpi: + comm = grid.distributor._obj_comm + else: + comm = 'PETSC_COMM_WORLD' + call_mpi = petsc_call_mpi('MPI_Comm_size', [comm, Byref(objs['size'])]) + # from IPython import embed; embed() + # get_size = c.Line('PetscCallMPI(MPI_Comm_size(comm,&(size)));') + rank = c.Line('PetscMPIInt rank;') + get_rank = c.Line('PetscCallMPI(MPI_Comm_rank(comm,&(rank)));') + print_comm_size = c.Line('PetscSynchronizedPrintf(comm, "MPI_Comm_size: %d\\n", size);') + + print_comm_rank = c.Line('PetscSynchronizedPrintf(comm, "MPI_Comm_rank: %d\\n", rank);') + + + flush = c.Line('PetscSynchronizedFlush(comm, PETSC_STDOUT);') + - return call_mpi, BlankLine + return call_mpi, rank, get_rank, print_comm_size, flush, print_comm_rank, flush, BlankLine + # return call_mpi, BlankLine class Builder: diff --git a/devito/petsc/iet/routines.py b/devito/petsc/iet/routines.py index ae305a5ba0..812e7b9b86 100644 --- a/devito/petsc/iet/routines.py +++ b/devito/petsc/iet/routines.py @@ -1086,6 +1086,8 @@ def __init__(self, **kwargs): self.objs = kwargs.get('objs') self.sregistry = kwargs.get('sregistry') self.grid = kwargs.get('grid') + # from IPython import embed; embed() + self.devito_mpi = kwargs['options'].get('mpi', False) self.fielddata = self.injectsolve.expr.rhs.fielddata self.solver_objs = self._build() @@ -1126,7 +1128,11 @@ def _build(self): 'callbackdm': CallbackDM(sreg.make_name(prefix='dm')), } # TODO: Devito MPI + PETSc testing - base_dict['comm'] = self.grid.distributor._obj_comm + # from IPython import embed; embed() + if self.devito_mpi: + base_dict['comm'] = self.grid.distributor._obj_comm + else: + base_dict['comm'] = 'PETSC_COMM_WORLD' self._target_dependent(base_dict) return self._extend_build(base_dict) diff --git a/examples/petsc/petsc_test.py b/examples/petsc/petsc_test.py index 5d93669d5f..eb74c639db 100644 --- a/examples/petsc/petsc_test.py +++ b/examples/petsc/petsc_test.py @@ -28,4 +28,4 @@ op = Operator(petsc) op.apply() -print(op.ccode) +# print(op.ccode) From 51d4d7e9dc91d03b26b15449c9fadae9ecc591b9 Mon Sep 17 00:00:00 2001 From: ZoeLeibowitz Date: Thu, 12 Jun 2025 00:19:13 +0100 Subject: [PATCH 48/84] misc: Address more comments and add docstrings --- devito/petsc/iet/routines.py | 33 ++++++++++++++++----------------- devito/petsc/types/types.py | 28 +++++++++++++++++++++++++--- devito/symbolics/extraction.py | 3 ++- 3 files changed, 43 insertions(+), 21 deletions(-) diff --git a/devito/petsc/iet/routines.py b/devito/petsc/iet/routines.py index 3ddf78e770..4efc88c1b5 100644 --- a/devito/petsc/iet/routines.py +++ b/devito/petsc/iet/routines.py @@ -258,13 +258,12 @@ def _create_matvec_body(self, body, jacobian): ) # Dereference function data in struct - dereference_funcs = [Dereference(i, ctx) for i in - fields if isinstance(i.function, AbstractFunction)] + derefs = self.dereference_funcs(ctx, fields) matvec_body = CallableBody( List(body=body), init=(objs['begin_user'],), - stacks=stacks+tuple(dereference_funcs), + stacks=stacks+derefs, retstmt=(Call('PetscFunctionReturn', arguments=[0]),) ) @@ -398,13 +397,12 @@ def _create_formfunc_body(self, body): ) # Dereference function data in struct - dereference_funcs = [Dereference(i, ctx) for i in - fields if isinstance(i.function, AbstractFunction)] + derefs = self.dereference_funcs(ctx, fields) formfunc_body = CallableBody( List(body=body), init=(objs['begin_user'],), - stacks=stacks+tuple(dereference_funcs), + stacks=stacks+derefs, retstmt=(Call('PetscFunctionReturn', arguments=[0]),)) # Replace non-function data with pointer to data in struct @@ -509,15 +507,12 @@ def _create_form_rhs_body(self, body): ) # Dereference function data in struct - dereference_funcs = tuple( - [Dereference(i, ctx) for i in - fields if isinstance(i.function, AbstractFunction)] - ) + derefs = self.dereference_funcs(ctx, fields) formrhs_body = CallableBody( List(body=[body]), init=(objs['begin_user'],), - stacks=stacks+dereference_funcs, + stacks=stacks+derefs, retstmt=(Call('PetscFunctionReturn', arguments=[0]),) ) @@ -590,13 +585,12 @@ def _create_initial_guess_body(self, body): ) # Dereference function data in struct - dereference_funcs = [Dereference(i, ctx) for i in - fields if isinstance(i.function, AbstractFunction)] + derefs = self.dereference_funcs(ctx, fields) body = CallableBody( List(body=[body]), init=(objs['begin_user'],), - stacks=stacks+tuple(dereference_funcs), + stacks=stacks+derefs, retstmt=(Call('PetscFunctionReturn', arguments=[0]),) ) @@ -662,6 +656,12 @@ def zero_vector(self, vec): """ return petsc_call('VecSet', [vec, 0.0]) if self.zero_memory else None + def dereference_funcs(self, struct, fields): + return tuple( + [Dereference(i, struct) for i in + fields if isinstance(i.function, AbstractFunction)] + ) + class CCBBuilder(CBBuilder): def __init__(self, **kwargs): @@ -894,8 +894,7 @@ def _whole_formfunc_body(self, body): ) # Dereference function data in struct - dereference_funcs = [Dereference(i, ctx) for i in - fields if isinstance(i.function, AbstractFunction)] + derefs = self.dereference_funcs(ctx, fields) f_soa = PointerCast(fbundle) x_soa = PointerCast(xbundle) @@ -903,7 +902,7 @@ def _whole_formfunc_body(self, body): formfunc_body = CallableBody( List(body=body), init=(objs['begin_user'],), - stacks=stacks+tuple(dereference_funcs), + stacks=stacks+derefs, casts=(f_soa, x_soa), retstmt=(Call('PetscFunctionReturn', arguments=[0]),) ) diff --git a/devito/petsc/types/types.py b/devito/petsc/types/types.py index 2e866896cb..eee16047cb 100644 --- a/devito/petsc/types/types.py +++ b/devito/petsc/types/types.py @@ -241,7 +241,7 @@ class Jacobian: This Jacobian is defined implicitly via matrix-vector products derived from the symbolic equations provided in `matvecs`. - It assumes the problem is linear, meaning the Jacobian + The class assumes the problem is linear, meaning the Jacobian corresponds to a constant coefficient matrix and does not require explicit symbolic differentiation. """ @@ -274,6 +274,7 @@ def _build_matvecs(self): matvecs.extend( e for e in self._build_matvec_eq(eq) if e is not None ) + matvecs = tuple(sorted(matvecs, key=lambda e: not isinstance(e, EssentialBC))) matvecs = self._scale_non_bcs(matvecs) @@ -469,6 +470,20 @@ def __repr__(self): class Residual: """ + Gennerates the metadata needed to define the nonlinear residual function + F(target) = 0 for use with PETSc's SNES interface. + + PETSc's SNES interface includes methods for solving nonlinear systems of + equations using Newton-type methods. For linear problems, `SNESKSPONLY` + is used to perform a single Newton iteration, unifying the + interface for both linear and nonlinear problems. + + This class encapsulates the symbolic equations used to construct the + residual function F(target) = F_(target) - b, where b contains all + terms independent of the solution `target`. + + References: + - https://petsc.org/main/manual/snes/ """ def __init__(self, target, eqns, arrays, time_mapper, scdiag): self.target = target @@ -481,12 +496,18 @@ def __init__(self, target, eqns, arrays, time_mapper, scdiag): @property def formfuncs(self): """ + Stores the equations used to build the `FormFunction` + callback generated at the IET level. This function is + passed to PETSc via `SNESSetFunction(..., FormFunction, ...)`. """ return self._formfuncs @property def formrhs(self): """ + Stores the equations used to generate the RHS + vector `b` through the `FormRHS` callback generated at the IET level. + The SNES solver is then called via `SNESSolve(..., b, target)`. """ return self._formrhs @@ -544,7 +565,7 @@ class MixedResidual(Residual): """ """ def __init__(self, target_eqns, arrays, time_mapper, scdiag): - self.targets = as_tuple(target_eqns.keys()) + self.targets = tuple(target_eqns.keys()) self.arrays = arrays self.time_mapper = time_mapper self.scdiag = scdiag @@ -592,6 +613,7 @@ def _build_function_eq(self, eq, target): self.arrays[target]['x'], eq.rhs, subdomain=eq.subdomain ) return (zero_row, zero_col) + else: if isinstance(zeroed, (int, float)): rhs = zeroed * volume @@ -599,7 +621,7 @@ def _build_function_eq(self, eq, target): rhs = zeroed.subs(mapper) rhs = rhs.subs(self.time_mapper)*volume - return as_tuple(Eq(self.arrays[target]['f'], rhs, subdomain=eq.subdomain)) + return (Eq(self.arrays[target]['f'], rhs, subdomain=eq.subdomain),) class InitialGuess: diff --git a/devito/symbolics/extraction.py b/devito/symbolics/extraction.py index fc47e94444..c0a28433cc 100644 --- a/devito/symbolics/extraction.py +++ b/devito/symbolics/extraction.py @@ -20,9 +20,10 @@ def separate_eqn(eqn, target): from devito.operations.solve import eval_time_derivatives zeroed_eqn = eval_time_derivatives(zeroed_eqn.lhs) - target_funcs = set(generate_targets(zeroed_eqn, target)) + target_funcs = set(generate_targets(zeroed_eqn, target)) b, F_target = remove_targets(zeroed_eqn, target_funcs) + return -b, F_target, zeroed_eqn, target_funcs From a4a56603fb92c4303df4b4c2f2568877875c552f Mon Sep 17 00:00:00 2001 From: ZoeLeibowitz Date: Thu, 12 Jun 2025 12:21:49 +0100 Subject: [PATCH 49/84] misc: Clean up more docstrings --- devito/petsc/iet/routines.py | 92 +++++++++---------- devito/petsc/solve.py | 91 +++++++++--------- devito/petsc/types/types.py | 172 +++++++++++++++++------------------ tests/test_petsc.py | 10 +- 4 files changed, 181 insertions(+), 184 deletions(-) diff --git a/devito/petsc/iet/routines.py b/devito/petsc/iet/routines.py index 4efc88c1b5..f762c3d30b 100644 --- a/devito/petsc/iet/routines.py +++ b/devito/petsc/iet/routines.py @@ -46,9 +46,9 @@ def __init__(self, **kwargs): self._user_struct_callback = None # TODO: Test pickling. The mutability of these lists # could cause issues when pickling? - self._matvecs = [] - self._formfuncs = [] - self._formrhs = [] + self._J_efuncs = [] + self._F_efuncs = [] + self._b_efuncs = [] self._initialguesses = [] self._make_core() @@ -73,23 +73,23 @@ def main_matvec_callback(self): is set in the main kernel via `PetscCall(MatShellSetOperation(J,MATOP_MULT,(void (*)(void))...));` """ - return self._matvecs[0] + return self._J_efuncs[0] @property def main_formfunc_callback(self): - return self._formfuncs[0] + return self._F_efuncs[0] @property - def matvecs(self): - return self._matvecs + def J_efuncs(self): + return self._J_efuncs @property - def formfuncs(self): - return self._formfuncs + def F_efuncs(self): + return self._F_efuncs @property - def formrhs(self): - return self._formrhs + def b_efuncs(self): + return self._b_efuncs @property def initialguesses(self): @@ -122,27 +122,27 @@ def _make_core(self): self._make_matvec(self.fielddata.jacobian) self._make_formfunc() self._make_formrhs() - if self.fielddata.initialguess.eqs: + if self.fielddata.initialguess.exprs: self._make_initialguess() self._make_user_struct_callback() def _make_matvec(self, jacobian, prefix='MatMult'): # Compile matvec `eqns` into an IET via recursive compilation matvecs = jacobian.matvecs - irs_matvec, _ = self.rcompile(matvecs, + irs, _ = self.rcompile(matvecs, options={'mpi': False}, sregistry=self.sregistry, concretize_mapper=self.concretize_mapper) - body_matvec = self._create_matvec_body(List(body=irs_matvec.uiet.body), + body = self._create_matvec_body(List(body=irs.uiet.body), jacobian) objs = self.objs cb = PETScCallable( self.sregistry.make_name(prefix=prefix), - body_matvec, + body, retval=objs['err'], parameters=(objs['J'], objs['X'], objs['Y']) ) - self._matvecs.append(cb) + self._J_efuncs.append(cb) self._efuncs[cb.name] = cb def _create_matvec_body(self, body, jacobian): @@ -260,7 +260,7 @@ def _create_matvec_body(self, body, jacobian): # Dereference function data in struct derefs = self.dereference_funcs(ctx, fields) - matvec_body = CallableBody( + body = CallableBody( List(body=body), init=(objs['begin_user'],), stacks=stacks+derefs, @@ -269,20 +269,20 @@ def _create_matvec_body(self, body, jacobian): # Replace non-function data with pointer to data in struct subs = {i._C_symbol: FieldFromPointer(i._C_symbol, ctx) for i in fields} - matvec_body = Uxreplace(subs).visit(matvec_body) + body = Uxreplace(subs).visit(body) self._struct_params.extend(fields) - return matvec_body + return body def _make_formfunc(self): - formfuncs = self.fielddata.residual.formfuncs - # Compile formfunc `eqns` into an IET via recursive compilation - irs_formfunc, _ = self.rcompile( - formfuncs, options={'mpi': False}, sregistry=self.sregistry, + F_exprs = self.fielddata.residual.F_exprs + # Compile `F_exprs` into an IET via recursive compilation + irs, _ = self.rcompile( + F_exprs, options={'mpi': False}, sregistry=self.sregistry, concretize_mapper=self.concretize_mapper ) body_formfunc = self._create_formfunc_body( - List(body=irs_formfunc.uiet.body) + List(body=irs.uiet.body) ) objs = self.objs cb = PETScCallable( @@ -291,7 +291,7 @@ def _make_formfunc(self): retval=objs['err'], parameters=(objs['snes'], objs['X'], objs['F'], objs['dummyptr']) ) - self._formfuncs.append(cb) + self._F_efuncs.append(cb) self._efuncs[cb.name] = cb def _create_formfunc_body(self, body): @@ -399,7 +399,7 @@ def _create_formfunc_body(self, body): # Dereference function data in struct derefs = self.dereference_funcs(ctx, fields) - formfunc_body = CallableBody( + body = CallableBody( List(body=body), init=(objs['begin_user'],), stacks=stacks+derefs, @@ -408,28 +408,28 @@ def _create_formfunc_body(self, body): # Replace non-function data with pointer to data in struct subs = {i._C_symbol: FieldFromPointer(i._C_symbol, ctx) for i in fields} - return Uxreplace(subs).visit(formfunc_body) + return Uxreplace(subs).visit(body) def _make_formrhs(self): - formrhs = self.fielddata.residual.formrhs + b_exprs = self.fielddata.residual.b_exprs sobjs = self.solver_objs - # Compile formrhs `eqns` into an IET via recursive compilation - irs_formrhs, _ = self.rcompile( - formrhs, options={'mpi': False}, sregistry=self.sregistry, + # Compile `b_exprs` into an IET via recursive compilation + irs, _ = self.rcompile( + b_exprs, options={'mpi': False}, sregistry=self.sregistry, concretize_mapper=self.concretize_mapper ) - body_formrhs = self._create_form_rhs_body( - List(body=irs_formrhs.uiet.body) + body = self._create_form_rhs_body( + List(body=irs.uiet.body) ) objs = self.objs cb = PETScCallable( self.sregistry.make_name(prefix='FormRHS'), - body_formrhs, + body, retval=objs['err'], parameters=(sobjs['callbackdm'], objs['B']) ) - self._formrhs.append(cb) + self._b_efuncs.append(cb) self._efuncs[cb.name] = cb def _create_form_rhs_body(self, body): @@ -509,7 +509,7 @@ def _create_form_rhs_body(self, body): # Dereference function data in struct derefs = self.dereference_funcs(ctx, fields) - formrhs_body = CallableBody( + body = CallableBody( List(body=[body]), init=(objs['begin_user'],), stacks=stacks+derefs, @@ -520,24 +520,24 @@ def _create_form_rhs_body(self, body): subs = {i._C_symbol: FieldFromPointer(i._C_symbol, ctx) for i in fields if not isinstance(i.function, AbstractFunction)} - return Uxreplace(subs).visit(formrhs_body) + return Uxreplace(subs).visit(body) def _make_initialguess(self): - initguess = self.fielddata.initialguess.eqs + exprs = self.fielddata.initialguess.exprs sobjs = self.solver_objs # Compile initital guess `eqns` into an IET via recursive compilation irs, _ = self.rcompile( - initguess, options={'mpi': False}, sregistry=self.sregistry, + exprs, options={'mpi': False}, sregistry=self.sregistry, concretize_mapper=self.concretize_mapper ) - body_init_guess = self._create_initial_guess_body( + body = self._create_initial_guess_body( List(body=irs.uiet.body) ) objs = self.objs cb = PETScCallable( self.sregistry.make_name(prefix='FormInitialGuess'), - body_init_guess, + body, retval=objs['err'], parameters=(sobjs['callbackdm'], objs['xloc']) ) @@ -770,10 +770,10 @@ def _whole_matvec_body(self): ) def _make_whole_formfunc(self): - formfuncs = self.fielddata.residual.formfuncs + F_exprs = self.fielddata.residual.F_exprs # Compile formfunc `eqns` into an IET via recursive compilation irs_formfunc, _ = self.rcompile( - formfuncs, options={'mpi': False}, sregistry=self.sregistry, + F_exprs, options={'mpi': False}, sregistry=self.sregistry, concretize_mapper=self.concretize_mapper ) body_formfunc = self._whole_formfunc_body(List(body=irs_formfunc.uiet.body)) @@ -1024,7 +1024,7 @@ def _submat_callback_body(self): iteration = Iteration(List(body=iter_body), i, upper_bound) nonzero_submats = self.jacobian.nonzero_submatrices - matvec_lookup = {mv.name.split('_')[0]: mv for mv in self.matvecs} + matvec_lookup = {mv.name.split('_')[0]: mv for mv in self.J_efuncs} matmult_op = [ petsc_call( @@ -1616,11 +1616,11 @@ def _execute_solve(self): struct_assignment = self.timedep.assign_time_iters(sobjs['userctx']) - rhs_callback = self.cbbuilder.formrhs[0] + b_efunc = self.cbbuilder.b_efuncs[0] dmda = sobjs['dmda'] - rhs_call = petsc_call(rhs_callback.name, [sobjs['dmda'], sobjs['bglobal']]) + rhs_call = petsc_call(b_efunc.name, [sobjs['dmda'], sobjs['bglobal']]) vec_place_array = self.timedep.place_array(target) diff --git a/devito/petsc/solve.py b/devito/petsc/solve.py index 9a85774ef8..949328cf67 100644 --- a/devito/petsc/solve.py +++ b/devito/petsc/solve.py @@ -12,9 +12,9 @@ __all__ = ['PETScSolve'] -def PETScSolve(target_eqns, target=None, solver_parameters=None): +def PETScSolve(target_exprs, target=None, solver_parameters=None): """ - Returns a symbolic equation representing a linear PETSc solver, + Returns a symbolic expression representing a linear PETSc solver, enriched with all the necessary metadata for execution within an `Operator`. When passed to an `Operator`, this symbolic equation triggers code generation and lowering to the PETSc backend. @@ -24,8 +24,8 @@ def PETScSolve(target_eqns, target=None, solver_parameters=None): Parameters ---------- - target_eqns : Eq or list of Eq, or dict of Function-like -> Eq or list of Eq - The targets and symbolic equations defining the system to be solved. + target_exprs : Eq or list of Eq, or dict of Function-like -> Eq or list of Eq + The targets and symbolic expressions defining the system to be solved. - Single-field problem: Pass a single Eq or list of Eq, and specify `target` separately: @@ -46,26 +46,26 @@ def PETScSolve(target_eqns, target=None, solver_parameters=None): Returns ------- - Eq - A symbolic equation that wraps the linear solver. + Eq: + A symbolic expression that wraps the linear solver. This can be passed directly to a Devito Operator. """ if target is not None: - return InjectSolve(solver_parameters, {target: target_eqns}).build_eq() + return InjectSolve(solver_parameters, {target: target_exprs}).build_expr() else: - return InjectMixedSolve(solver_parameters, target_eqns).build_eq() + return InjectMixedSolve(solver_parameters, target_exprs).build_expr() class InjectSolve: - def __init__(self, solver_parameters=None, target_eqns=None): + def __init__(self, solver_parameters=None, target_exprs=None): self.solver_params = solver_parameters self.time_mapper = None - self.target_eqns = target_eqns + self.target_exprs = target_exprs - def build_eq(self): + def build_expr(self): target, funcs, fielddata = self.linear_solve_args() - # Placeholder equation for inserting calls to the solver + # Placeholder expression for inserting calls to the solver linear_solve = LinearSolveExpr( funcs, self.solver_params, @@ -76,18 +76,18 @@ def build_eq(self): return [PetscEq(target, linear_solve)] def linear_solve_args(self): - target, eqns = next(iter(self.target_eqns.items())) - eqns = as_tuple(eqns) + target, exprs = next(iter(self.target_exprs.items())) + exprs = as_tuple(exprs) - funcs = get_funcs(eqns) + funcs = get_funcs(exprs) self.time_mapper = generate_time_mapper(funcs) - arrays = self.generate_arrays_combined(target) + arrays = self.generate_arrays(target) - eqns = sorted(eqns, key=lambda e: not isinstance(e, EssentialBC)) + exprs = sorted(exprs, key=lambda e: not isinstance(e, EssentialBC)) - jacobian = Jacobian(target, eqns, arrays, self.time_mapper) - residual = Residual(target, eqns, arrays, self.time_mapper, jacobian.scdiag) - initialguess = InitialGuess(target, eqns, arrays) + jacobian = Jacobian(target, exprs, arrays, self.time_mapper) + residual = Residual(target, exprs, arrays, self.time_mapper, jacobian.scdiag) + initialguess = InitialGuess(target, exprs, arrays) field_data = FieldData( target=target, @@ -98,55 +98,55 @@ def linear_solve_args(self): ) return target, tuple(funcs), field_data - - def generate_arrays(self, target): + + def generate_arrays(self, *targets): return { - p: PETScArray(name=f'{p}_{target.name}', - target=target, - liveness='eager', - localinfo=localinfo) - for p in prefixes + t: { + p: PETScArray(name=f'{p}_{t.name}', + target=t, + liveness='eager', + localinfo=localinfo) + for p in prefixes + } + for t in targets } - def generate_arrays_combined(self, *targets): - return {target: self.generate_arrays(target) for target in targets} - class InjectMixedSolve(InjectSolve): def linear_solve_args(self): - combined_eqns = [] - for eqns in self.target_eqns.values(): - combined_eqns.extend(eqns) - funcs = get_funcs(combined_eqns) - self.time_mapper = generate_time_mapper(funcs) + exprs = [] + for e in self.target_exprs.values(): + exprs.extend(e) - coupled_targets = list(self.target_eqns.keys()) + funcs = get_funcs(exprs) + self.time_mapper = generate_time_mapper(funcs) - arrays = self.generate_arrays_combined(*coupled_targets) + targets = list(self.target_exprs.keys()) + arrays = self.generate_arrays(*targets) jacobian = MixedJacobian( - self.target_eqns, arrays, self.time_mapper + self.target_exprs, arrays, self.time_mapper ) residual = MixedResidual( - self.target_eqns, arrays, self.time_mapper, + self.target_exprs, arrays, self.time_mapper, jacobian.target_scaler_mapper ) all_data = MultipleFieldData( - targets=coupled_targets, + targets=targets, arrays=arrays, jacobian=jacobian, residual=residual ) - return coupled_targets[0], tuple(funcs), all_data + return targets[0], tuple(funcs), all_data def generate_time_mapper(funcs): """ - Replace time indices with `Symbols` in equations used within + Replace time indices with `Symbols` in expressions used within PETSc callback functions. These symbols are Uxreplaced at the IET level to align with the `TimeDimension` and `ModuloDimension` objects present in the initial lowering. @@ -176,11 +176,10 @@ def generate_time_mapper(funcs): return dict(zip(time_indices, tau_symbs)) -def get_funcs(eqns): +def get_funcs(exprs): funcs = [ - func - for eq in eqns - for func in retrieve_functions(eval_time_derivatives(eq.lhs - eq.rhs)) + f for e in exprs + for f in retrieve_functions(eval_time_derivatives(e.lhs - e.rhs)) ] return filter_ordered(funcs) diff --git a/devito/petsc/types/types.py b/devito/petsc/types/types.py index eee16047cb..f4e92f2f2e 100644 --- a/devito/petsc/types/types.py +++ b/devito/petsc/types/types.py @@ -239,21 +239,26 @@ class Jacobian: Represents a Jacobian matrix. This Jacobian is defined implicitly via matrix-vector products - derived from the symbolic equations provided in `matvecs`. + derived from the symbolic expressions provided in `matvecs`. The class assumes the problem is linear, meaning the Jacobian corresponds to a constant coefficient matrix and does not require explicit symbolic differentiation. """ - def __init__(self, target, eqns, arrays, time_mapper): + def __init__(self, target, exprs, arrays, time_mapper): self.target = target - self.eqns = eqns + self.exprs = exprs self.arrays = arrays self.time_mapper = time_mapper self._build_matvecs() @property def matvecs(self): + """ + Stores the expressions used to generate the `MatMult` + callback generated at the IET level. This function is + passed to PETSc via `MatShellSetOperation(...,MATOP_MULT,(void (*)(void))MatMult)`. + """ return self._matvecs @property @@ -270,9 +275,9 @@ def col_target(self): def _build_matvecs(self): matvecs = [] - for eq in self.eqns: + for eq in self.exprs: matvecs.extend( - e for e in self._build_matvec_eq(eq) if e is not None + e for e in self._build_matvec_expr(eq) if e is not None ) matvecs = tuple(sorted(matvecs, key=lambda e: not isinstance(e, EssentialBC))) @@ -284,32 +289,32 @@ def _build_matvecs(self): self._matvecs = matvecs self._scdiag = scdiag - def _build_matvec_eq(self, eq, col_target=None, row_target=None): + def _build_matvec_expr(self, expr, col_target=None, row_target=None): col_target = col_target or self.target row_target = row_target or self.target - _, F_target, _, targets = separate_eqn(eq, col_target) + _, F_target, _, targets = separate_eqn(expr, col_target) if F_target: return self._make_matvec( - eq, F_target, targets, col_target, row_target + expr, F_target, targets, col_target, row_target ) else: return (None,) - def _make_matvec(self, eq, F_target, targets, col_target, row_target): + def _make_matvec(self, expr, F_target, targets, col_target, row_target): y = self.arrays[row_target]['y'] x = self.arrays[col_target]['x'] - if isinstance(eq, EssentialBC): + if isinstance(expr, EssentialBC): # NOTE: Essential BCs are trivial equations in the solver. # See `EssentialBC` for more details. - zero_row = ZeroRow(y, x, subdomain=eq.subdomain) - zero_column = ZeroColumn(x, 0., subdomain=eq.subdomain) + zero_row = ZeroRow(y, x, subdomain=expr.subdomain) + zero_column = ZeroColumn(x, 0., subdomain=expr.subdomain) return (zero_row, zero_column) else: rhs = F_target.subs(targets_to_arrays(x, targets)) rhs = rhs.subs(self.time_mapper) - return (Eq(y, rhs, subdomain=eq.subdomain),) + return (Eq(y, rhs, subdomain=expr.subdomain),) def _scale_non_bcs(self, matvecs, target=None): target = target or self.target @@ -376,8 +381,6 @@ class MixedJacobian(Jacobian): # TODO: pcfieldsplit support for each block """ def __init__(self, target_eqns, arrays, time_mapper): - """ - """ self.targets = tuple(target_eqns.keys()) self.arrays = arrays self.time_mapper = time_mapper @@ -394,15 +397,13 @@ def submatrices(self): @property def n_submatrices(self): - """ - Return the number of submatrix blocks. - """ + """Return the number of submatrix blocks.""" return len(self._submatrices) @property def nonzero_submatrices(self): """Return SubMatrixBlock objects that have non-empty matvecs.""" - return [submat for submat in self.submatrices if submat.matvecs] + return [m for m in self.submatrices if m.matvecs] @property def target_scaler_mapper(self): @@ -411,22 +412,22 @@ def target_scaler_mapper(self): diagonal subblock. """ mapper = {} - for sm in self.submatrices: - if sm.row_idx == sm.col_idx: - mapper[sm.row_target] = sm.scdiag + for m in self.submatrices: + if m.row_idx == m.col_idx: + mapper[m.row_target] = m.scdiag return mapper - def _build_blocks(self, target_eqns): + def _build_blocks(self, target_exprs): """ Build all SubMatrixBlock objects for the Jacobian. """ for i, row_target in enumerate(self.targets): - eqns = target_eqns[row_target] + exprs = target_exprs[row_target] for j, col_target in enumerate(self.targets): matvecs = [] - for eq in eqns: + for expr in exprs: matvecs.extend( - e for e in self._build_matvec_eq(eq, col_target, row_target) + e for e in self._build_matvec_expr(expr, col_target, row_target) ) matvecs = [m for m in matvecs if m is not None] @@ -475,56 +476,56 @@ class Residual: PETSc's SNES interface includes methods for solving nonlinear systems of equations using Newton-type methods. For linear problems, `SNESKSPONLY` - is used to perform a single Newton iteration, unifying the + can be used to perform a single Newton iteration, unifying the interface for both linear and nonlinear problems. - This class encapsulates the symbolic equations used to construct the + This class encapsulates the symbolic expressions used to construct the residual function F(target) = F_(target) - b, where b contains all terms independent of the solution `target`. References: - https://petsc.org/main/manual/snes/ """ - def __init__(self, target, eqns, arrays, time_mapper, scdiag): + def __init__(self, target, exprs, arrays, time_mapper, scdiag): self.target = target - self.eqns = eqns + self.exprs = exprs self.arrays = arrays self.time_mapper = time_mapper self.scdiag = scdiag - self._build_equations() + self._build_exprs() @property - def formfuncs(self): + def F_exprs(self): """ - Stores the equations used to build the `FormFunction` + Stores the expressions used to build the `FormFunction` callback generated at the IET level. This function is passed to PETSc via `SNESSetFunction(..., FormFunction, ...)`. """ - return self._formfuncs + return self._F_exprs @property - def formrhs(self): + def b_exprs(self): """ - Stores the equations used to generate the RHS + Stores the expressions used to generate the RHS vector `b` through the `FormRHS` callback generated at the IET level. The SNES solver is then called via `SNESSolve(..., b, target)`. """ - return self._formrhs + return self._b_exprs - def _build_equations(self): + def _build_exprs(self): """ """ - funcs = [] - rhs = [] + F_exprs = [] + b_exprs = [] - for eq in self.eqns: - b, F_target, _, targets = separate_eqn(eq, self.target) - funcs.extend(self._make_F_target(eq, F_target, targets)) + for e in self.exprs: + b, F_target, _, targets = separate_eqn(e, self.target) + F_exprs.extend(self._make_F_target(e, F_target, targets)) # TODO: If b is zero then don't need a rhs vector+callback - rhs.extend(self._make_b(eq, b)) + b_exprs.extend(self._make_b(e, b)) - self._formfuncs = tuple([self._scale_bcs(eq) for eq in funcs]) - self._formrhs = tuple(rhs) + self._F_exprs = tuple([self._scale_bcs(e) for e in F_exprs]) + self._b_exprs = tuple(b_exprs) def _make_F_target(self, eq, F_target, targets): arrays = self.arrays[self.target] @@ -547,15 +548,15 @@ def _make_F_target(self, eq, F_target, targets): rhs = rhs.subs(self.time_mapper) * volume return (Eq(arrays['f'], rhs, subdomain=eq.subdomain),) - def _make_b(self, eq, b): + def _make_b(self, expr, b): b_arr = self.arrays[self.target]['b'] - rhs = 0. if isinstance(eq, EssentialBC) else b.subs(self.time_mapper) + rhs = 0. if isinstance(expr, EssentialBC) else b.subs(self.time_mapper) rhs = rhs * self.target.grid.symbolic_volume_cell - return (Eq(b_arr, rhs, subdomain=eq.subdomain),) + return (Eq(b_arr, rhs, subdomain=expr.subdomain),) def _scale_bcs(self, eq, scdiag=None): """ - Scale ZeroRow equations using scdiag + Scale ZeroRow exprs using scdiag """ scdiag = scdiag or self.scdiag return eq._rebuild(rhs=scdiag * eq.rhs) if isinstance(eq, ZeroRow) else eq @@ -564,37 +565,35 @@ def _scale_bcs(self, eq, scdiag=None): class MixedResidual(Residual): """ """ - def __init__(self, target_eqns, arrays, time_mapper, scdiag): - self.targets = tuple(target_eqns.keys()) + def __init__(self, target_exprs, arrays, time_mapper, scdiag): + self.targets = tuple(target_exprs.keys()) self.arrays = arrays self.time_mapper = time_mapper self.scdiag = scdiag - self._build_equations(target_eqns) + self._build_exprs(target_exprs) @property - def formrhs(self): + def b_exprs(self): """ """ return None - def _build_equations(self, target_eqns): - all_formfuncs = [] - for target, eqns in target_eqns.items(): + def _build_exprs(self, target_exprs): + residual_exprs = [] + for t, exprs in target_exprs.items(): - formfuncs = chain.from_iterable( - self._build_function_eq(eq, target) - for eq in as_tuple(eqns) - ) - all_formfuncs.extend(formfuncs) + residual_exprs.extend( + chain.from_iterable(self._build_residual(e, t) + for e in as_tuple(exprs) + )) - self._formfuncs = tuple(sorted( - all_formfuncs, key=lambda e: not isinstance(e, EssentialBC) + self._F_exprs = tuple(sorted( + residual_exprs, key=lambda e: not isinstance(e, EssentialBC) )) - def _build_function_eq(self, eq, target): - zeroed = eq.lhs - eq.rhs - - zeroed_eqn = Eq(eq.lhs - eq.rhs, 0) + def _build_residual(self, expr, target): + zeroed = expr.lhs - expr.rhs + zeroed_eqn = Eq(zeroed, 0) eval_zeroed_eqn = eval_time_derivatives(zeroed_eqn.lhs) volume = target.grid.symbolic_volume_cell @@ -604,13 +603,13 @@ def _build_function_eq(self, eq, target): target_funcs = set(generate_targets(Eq(eval_zeroed_eqn, 0), t)) mapper.update(targets_to_arrays(self.arrays[t]['x'], target_funcs)) - if isinstance(eq, EssentialBC): - rhs = (self.arrays[target]['x'] - eq.rhs)*self.scdiag[target] + if isinstance(expr, EssentialBC): + rhs = (self.arrays[target]['x'] - expr.rhs)*self.scdiag[target] zero_row = ZeroRow( - self.arrays[target]['f'], rhs, subdomain=eq.subdomain + self.arrays[target]['f'], rhs, subdomain=expr.subdomain ) zero_col = ZeroColumn( - self.arrays[target]['x'], eq.rhs, subdomain=eq.subdomain + self.arrays[target]['x'], expr.rhs, subdomain=expr.subdomain ) return (zero_row, zero_col) @@ -621,41 +620,40 @@ def _build_function_eq(self, eq, target): rhs = zeroed.subs(mapper) rhs = rhs.subs(self.time_mapper)*volume - return (Eq(self.arrays[target]['f'], rhs, subdomain=eq.subdomain),) + return (Eq(self.arrays[target]['f'], rhs, subdomain=expr.subdomain),) class InitialGuess: """ Enforce initial guess to satisfy essential BCs. """ - def __init__(self, target, eqns, arrays): + def __init__(self, target, exprs, arrays): self.target = target - self.eqns = as_tuple(eqns) self.arrays = arrays - self._build_equations() + self._build_exprs(as_tuple(exprs)) @property - def eqs(self): + def exprs(self): """ """ - return self._eqs + return self._exprs - def _build_equations(self): + def _build_exprs(self, exprs): """ - Return a list of initial guess equations. + Return a list of initial guess expressions. """ - self._eqs = tuple([ + self._exprs = tuple([ eq for eq in - (self._make_initial_guess(e) for e in self.eqns) + (self._make_initial_guess(e) for e in exprs) if eq is not None ]) - def _make_initial_guess(self, eq): - if isinstance(eq, EssentialBC): - assert eq.lhs == self.target + def _make_initial_guess(self, expr): + if isinstance(expr, EssentialBC): + assert expr.lhs == self.target return Eq( - self.arrays[self.target]['x'], eq.rhs, - subdomain=eq.subdomain + self.arrays[self.target]['x'], expr.rhs, + subdomain=expr.subdomain ) else: return None diff --git a/tests/test_petsc.py b/tests/test_petsc.py index 10e030b064..c5f9c42b43 100644 --- a/tests/test_petsc.py +++ b/tests/test_petsc.py @@ -100,12 +100,12 @@ 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 == 'MatMult0'] + matvec_efunc = [root for root in callable_roots if root.name == 'MatMult0'] - formrhs_callback = [root for root in callable_roots if root.name == 'FormRHS0'] + b_efunc = [root for root in callable_roots if root.name == 'FormRHS0'] - action_expr = FindNodes(Expression).visit(matvec_callback[0]) - rhs_expr = FindNodes(Expression).visit(formrhs_callback[0]) + action_expr = FindNodes(Expression).visit(matvec_efunc[0]) + rhs_expr = FindNodes(Expression).visit(b_efunc[0]) assert str(action_expr[-1].expr.rhs) == ( '(x_f[x + 1, y + 2]/ctx0->h_x**2' @@ -127,7 +127,7 @@ def test_petsc_solve(): assert len(retrieve_iteration_tree(op)) == 0 # TODO: Remove pragmas from PETSc callback functions - assert len(matvec_callback[0].parameters) == 3 + assert len(matvec_efunc[0].parameters) == 3 @skipif('petsc') From 312c8266aa14468ccabb0f633a6954dce1c8c2ad Mon Sep 17 00:00:00 2001 From: ZoeLeibowitz Date: Thu, 12 Jun 2025 17:18:41 +0100 Subject: [PATCH 50/84] misc: Add BaseJacobian --- devito/petsc/types/types.py | 216 ++++++++++++++++++++++-------------- 1 file changed, 131 insertions(+), 85 deletions(-) diff --git a/devito/petsc/types/types.py b/devito/petsc/types/types.py index f4e92f2f2e..5e9796a5d8 100644 --- a/devito/petsc/types/types.py +++ b/devito/petsc/types/types.py @@ -137,6 +137,26 @@ def eval(cls, *args): class FieldData: + """ + Metadata class passed to `LinearSolveExpr`. Encapsulates metadata for a single + `target` field needed to interface with PETSc SNES solvers. + + Parameters + ---------- + + target : Function-like + The target field to solve into, which is a Function-like object. + jacobian : Jacobian + Defines the matrix-vector product for the linear system, where the vector is + the PETScArray representing the `target`. + residual : Residual + Defines the nonlinear residual function F(target) = 0. + initialguess : InitialGuess + Defines the initial guess for the solution, which satisfies + essential boundary conditions. + arrays : dict + A dictionary mapping `target` to its corresponding PETScArrays. + """ def __init__(self, target=None, jacobian=None, residual=None, initialguess=None, arrays=None, **kwargs): self._target = target @@ -190,6 +210,24 @@ def targets(self): class MultipleFieldData(FieldData): + """ + Metadata class passed to `LinearSolveExpr`, for mixed-field problems, + where the solution vector spans multiple `targets`. + + Parameters + ---------- + targets : list of Function-like + The fields to solve into, each represented by a Function-like object. + jacobian : MixedJacobian + Defines the matrix-vector products for the full system Jacobian. + residual : MixedResidual + Defines the nonlinear residual function F(targets) = 0. + initialguess : InitialGuess + Defines the initial guess metadata, which satisfies + essential boundary conditions. + arrays : dict + A dictionary mapping the `targets` to their corresponding PETScArrays. + """ def __init__(self, targets, arrays, jacobian=None, residual=None): self._targets = as_tuple(targets) self._arrays = arrays @@ -208,6 +246,7 @@ def space_dimensions(self): @property def grid(self): + """The unique `Grid` associated with all targets.""" grids = [t.grid for t in self.targets] if len(set(grids)) > 1: raise ValueError( @@ -233,8 +272,70 @@ def space_order(self): def targets(self): return self._targets + +class BaseJacobian: + def __init__(self, arrays, target=None): + self.arrays = arrays + self.target = target + + def _scale_non_bcs(self, matvecs, target=None): + target = target or self.target + vol = target.grid.symbolic_volume_cell + + return [ + m if isinstance(m, EssentialBC) else m._rebuild(rhs=m.rhs * vol) + for m in matvecs + ] + + def _compute_scdiag(self, matvecs, col_target=None): + """ + """ + x = self.arrays[col_target or self.target]['x'] + + centres = { + centre_stencil(m.rhs, x, as_coeff=True) + for m in matvecs if not isinstance(m, EssentialBC) + } + return centres.pop() if len(centres) == 1 else 1.0 + + def _scale_bcs(self, matvecs, scdiag): + """ + Scale the essential BCs + """ + return [ + m._rebuild(rhs=m.rhs * scdiag) if isinstance(m, ZeroRow) else m + for m in matvecs + ] + + def _build_matvec_expr(self, expr, **kwargs): + col_target = kwargs.get('col_target', self.target) + row_target = kwargs.get('row_target', self.target) + + _, F_target, _, targets = separate_eqn(expr, col_target) + if F_target: + return self._make_matvec( + expr, F_target, targets, col_target, row_target + ) + else: + return (None,) + + def _make_matvec(self, expr, F_target, targets, col_target, row_target): + y = self.arrays[row_target]['y'] + x = self.arrays[col_target]['x'] + + if isinstance(expr, EssentialBC): + # NOTE: Essential BCs are trivial equations in the solver. + # See `EssentialBC` for more details. + zero_row = ZeroRow(y, x, subdomain=expr.subdomain) + zero_column = ZeroColumn(x, 0., subdomain=expr.subdomain) + return (zero_row, zero_column) + else: + rhs = F_target.subs(targets_to_arrays(x, targets)) + rhs = rhs.subs(self.time_mapper) + return (Eq(y, rhs, subdomain=expr.subdomain),) + -class Jacobian: +class Jacobian(BaseJacobian): """ Represents a Jacobian matrix. @@ -246,14 +347,14 @@ class Jacobian: require explicit symbolic differentiation. """ def __init__(self, target, exprs, arrays, time_mapper): - self.target = target + super().__init__(arrays=arrays, target=target) self.exprs = exprs - self.arrays = arrays self.time_mapper = time_mapper self._build_matvecs() @property def matvecs(self): + # TODO: add shortcut explanation etc """ Stores the expressions used to generate the `MatMult` callback generated at the IET level. This function is @@ -279,7 +380,6 @@ def _build_matvecs(self): matvecs.extend( e for e in self._build_matvec_expr(eq) if e is not None ) - matvecs = tuple(sorted(matvecs, key=lambda e: not isinstance(e, EssentialBC))) matvecs = self._scale_non_bcs(matvecs) @@ -289,83 +389,8 @@ def _build_matvecs(self): self._matvecs = matvecs self._scdiag = scdiag - def _build_matvec_expr(self, expr, col_target=None, row_target=None): - col_target = col_target or self.target - row_target = row_target or self.target - _, F_target, _, targets = separate_eqn(expr, col_target) - if F_target: - return self._make_matvec( - expr, F_target, targets, col_target, row_target - ) - else: - return (None,) - - def _make_matvec(self, expr, F_target, targets, col_target, row_target): - y = self.arrays[row_target]['y'] - x = self.arrays[col_target]['x'] - - if isinstance(expr, EssentialBC): - # NOTE: Essential BCs are trivial equations in the solver. - # See `EssentialBC` for more details. - zero_row = ZeroRow(y, x, subdomain=expr.subdomain) - zero_column = ZeroColumn(x, 0., subdomain=expr.subdomain) - return (zero_row, zero_column) - else: - rhs = F_target.subs(targets_to_arrays(x, targets)) - rhs = rhs.subs(self.time_mapper) - return (Eq(y, rhs, subdomain=expr.subdomain),) - - def _scale_non_bcs(self, matvecs, target=None): - target = target or self.target - vol = target.grid.symbolic_volume_cell - - return [ - m if isinstance(m, EssentialBC) else m._rebuild(rhs=m.rhs * vol) - for m in matvecs - ] - - def _compute_scdiag(self, matvecs, col_target=None): - """ - """ - x = self.arrays[col_target or self.target]['x'] - - centres = { - centre_stencil(m.rhs, x, as_coeff=True) - for m in matvecs if not isinstance(m, EssentialBC) - } - return centres.pop() if len(centres) == 1 else 1.0 - - def _scale_bcs(self, matvecs, scdiag): - """ - Scale the essential BCs - """ - return [ - m._rebuild(rhs=m.rhs * scdiag) if isinstance(m, ZeroRow) else m - for m in matvecs - ] - - -class SubMatrixBlock: - def __init__(self, name, matvecs, scdiag, row_target, - col_target, row_idx, col_idx, linear_idx): - self.name = name - self.matvecs = matvecs - self.scdiag = scdiag - self.row_target = row_target - self.col_target = col_target - self.row_idx = row_idx - self.col_idx = col_idx - self.linear_idx = linear_idx - - def is_diag(self): - return self.row_idx == self.col_idx - - def __repr__(self): - return (f"") - - -class MixedJacobian(Jacobian): +class MixedJacobian(BaseJacobian): """ Represents a Jacobian for a linear system with a solution vector composed of multiple fields (targets). @@ -380,12 +405,12 @@ class MixedJacobian(Jacobian): # TODO: pcfieldsplit support for each block """ - def __init__(self, target_eqns, arrays, time_mapper): - self.targets = tuple(target_eqns.keys()) - self.arrays = arrays + def __init__(self, target_exprs, arrays, time_mapper): + super().__init__(arrays=arrays, target=None) + self.targets = tuple(target_exprs.keys()) self.time_mapper = time_mapper self._submatrices = [] - self._build_blocks(target_eqns) + self._build_blocks(target_exprs) @property def submatrices(self): @@ -427,7 +452,9 @@ def _build_blocks(self, target_exprs): matvecs = [] for expr in exprs: matvecs.extend( - e for e in self._build_matvec_expr(expr, col_target, row_target) + e for e in self._build_matvec_expr( + expr, col_target=col_target, row_target=row_target + ) ) matvecs = [m for m in matvecs if m is not None] @@ -469,6 +496,25 @@ def __repr__(self): return f"" +class SubMatrixBlock: + def __init__(self, name, matvecs, scdiag, row_target, + col_target, row_idx, col_idx, linear_idx): + self.name = name + self.matvecs = matvecs + self.scdiag = scdiag + self.row_target = row_target + self.col_target = col_target + self.row_idx = row_idx + self.col_idx = col_idx + self.linear_idx = linear_idx + + def is_diag(self): + return self.row_idx == self.col_idx + + def __repr__(self): + return (f"") + + class Residual: """ Gennerates the metadata needed to define the nonlinear residual function From 048f693472d9b2b59673c33b93579592105cbc62 Mon Sep 17 00:00:00 2001 From: ZoeLeibowitz Date: Fri, 13 Jun 2025 14:35:34 +0100 Subject: [PATCH 51/84] misc: Docstrings, stop list output for PETScSolve, tests --- devito/petsc/iet/routines.py | 12 ++-- devito/petsc/solve.py | 14 ++-- devito/petsc/types/types.py | 70 ++++++++++-------- examples/petsc/cfd/01_navierstokes.py | 2 +- .../petsc/seismic/01_staggered_acoustic.py | 4 +- tests/test_petsc.py | 71 +++++++++++++++---- 6 files changed, 119 insertions(+), 54 deletions(-) diff --git a/devito/petsc/iet/routines.py b/devito/petsc/iet/routines.py index f762c3d30b..f87eaab78b 100644 --- a/devito/petsc/iet/routines.py +++ b/devito/petsc/iet/routines.py @@ -129,11 +129,13 @@ def _make_core(self): def _make_matvec(self, jacobian, prefix='MatMult'): # Compile matvec `eqns` into an IET via recursive compilation matvecs = jacobian.matvecs - irs, _ = self.rcompile(matvecs, - options={'mpi': False}, sregistry=self.sregistry, - concretize_mapper=self.concretize_mapper) - body = self._create_matvec_body(List(body=irs.uiet.body), - jacobian) + irs, _ = self.rcompile( + matvecs, options={'mpi': False}, sregistry=self.sregistry, + concretize_mapper=self.concretize_mapper + ) + body = self._create_matvec_body( + List(body=irs.uiet.body), jacobian + ) objs = self.objs cb = PETScCallable( diff --git a/devito/petsc/solve.py b/devito/petsc/solve.py index 949328cf67..1f4423707e 100644 --- a/devito/petsc/solve.py +++ b/devito/petsc/solve.py @@ -73,7 +73,7 @@ def build_expr(self): time_mapper=self.time_mapper, localinfo=localinfo ) - return [PetscEq(target, linear_solve)] + return PetscEq(target, linear_solve) def linear_solve_args(self): target, exprs = next(iter(self.target_exprs.items())) @@ -98,14 +98,16 @@ def linear_solve_args(self): ) return target, tuple(funcs), field_data - + def generate_arrays(self, *targets): return { t: { - p: PETScArray(name=f'{p}_{t.name}', - target=t, - liveness='eager', - localinfo=localinfo) + p: PETScArray( + name=f'{p}_{t.name}', + target=t, + liveness='eager', + localinfo=localinfo + ) for p in prefixes } for t in targets diff --git a/devito/petsc/types/types.py b/devito/petsc/types/types.py index 5e9796a5d8..b97b6ed406 100644 --- a/devito/petsc/types/types.py +++ b/devito/petsc/types/types.py @@ -138,8 +138,8 @@ def eval(cls, *args): class FieldData: """ - Metadata class passed to `LinearSolveExpr`. Encapsulates metadata for a single - `target` field needed to interface with PETSc SNES solvers. + Metadata for a single `target` field passed to `LinearSolveExpr`. + Used to interface with PETSc SNES solvers at the IET level. Parameters ---------- @@ -147,10 +147,10 @@ class FieldData: target : Function-like The target field to solve into, which is a Function-like object. jacobian : Jacobian - Defines the matrix-vector product for the linear system, where the vector is + Defines the matrix-vector product for the linear system, where the vector is the PETScArray representing the `target`. residual : Residual - Defines the nonlinear residual function F(target) = 0. + Defines the residual function F(target) = 0. initialguess : InitialGuess Defines the initial guess for the solution, which satisfies essential boundary conditions. @@ -212,7 +212,8 @@ def targets(self): class MultipleFieldData(FieldData): """ Metadata class passed to `LinearSolveExpr`, for mixed-field problems, - where the solution vector spans multiple `targets`. + where the solution vector spans multiple `targets`. Used to interface + with PETSc SNES solvers at the IET level. Parameters ---------- @@ -221,7 +222,7 @@ class MultipleFieldData(FieldData): jacobian : MixedJacobian Defines the matrix-vector products for the full system Jacobian. residual : MixedResidual - Defines the nonlinear residual function F(targets) = 0. + Defines the residual function F(targets) = 0. initialguess : InitialGuess Defines the initial guess metadata, which satisfies essential boundary conditions. @@ -272,13 +273,17 @@ def space_order(self): def targets(self): return self._targets - + class BaseJacobian: def __init__(self, arrays, target=None): self.arrays = arrays self.target = target def _scale_non_bcs(self, matvecs, target=None): + """ + Scale the symbolic expressions `matvecs` by the grid cell volume, + excluding EssentialBCs. + """ target = target or self.target vol = target.grid.symbolic_volume_cell @@ -287,8 +292,20 @@ def _scale_non_bcs(self, matvecs, target=None): for m in matvecs ] + def _scale_bcs(self, matvecs, scdiag): + """ + Scale the EssentialBCs in `matvecs` by `scdiag`. + """ + return [ + m._rebuild(rhs=m.rhs * scdiag) if isinstance(m, ZeroRow) else m + for m in matvecs + ] + def _compute_scdiag(self, matvecs, col_target=None): """ + Compute the diagonal scaling factor from the symbolic matrix-vector + expressions in `matvecs`. If the centre stencil (i.e. the diagonal term of the + matrix) is not unique, defaults to 1.0. """ x = self.arrays[col_target or self.target]['x'] @@ -298,15 +315,6 @@ def _compute_scdiag(self, matvecs, col_target=None): } return centres.pop() if len(centres) == 1 else 1.0 - def _scale_bcs(self, matvecs, scdiag): - """ - Scale the essential BCs - """ - return [ - m._rebuild(rhs=m.rhs * scdiag) if isinstance(m, ZeroRow) else m - for m in matvecs - ] - def _build_matvec_expr(self, expr, **kwargs): col_target = kwargs.get('col_target', self.target) row_target = kwargs.get('row_target', self.target) @@ -356,9 +364,9 @@ def __init__(self, target, exprs, arrays, time_mapper): def matvecs(self): # TODO: add shortcut explanation etc """ - Stores the expressions used to generate the `MatMult` - callback generated at the IET level. This function is - passed to PETSc via `MatShellSetOperation(...,MATOP_MULT,(void (*)(void))MatMult)`. + Stores the expressions used to generate the `MatMult` callback generated + at the IET level. This function is passed to PETSc via + `MatShellSetOperation(...,MATOP_MULT,(void (*)(void))MatMult)`. """ return self._matvecs @@ -517,7 +525,7 @@ def __repr__(self): class Residual: """ - Gennerates the metadata needed to define the nonlinear residual function + Generates the metadata needed to define the nonlinear residual function F(target) = 0 for use with PETSc's SNES interface. PETSc's SNES interface includes methods for solving nonlinear systems of @@ -610,6 +618,8 @@ def _scale_bcs(self, eq, scdiag=None): class MixedResidual(Residual): """ + Generates the metadata needed to define the nonlinear residual function + F(targets) = 0 for use with PETSc's SNES interface. """ def __init__(self, target_exprs, arrays, time_mapper, scdiag): self.targets = tuple(target_exprs.keys()) @@ -621,6 +631,9 @@ def __init__(self, target_exprs, arrays, time_mapper, scdiag): @property def b_exprs(self): """ + For mixed solvers, a callback to form the RHS vector `b` is not generated, + a single residual callback is generated to compute F(targets). + TODO: Investigate if this is optimal. """ return None @@ -629,10 +642,10 @@ def _build_exprs(self, target_exprs): for t, exprs in target_exprs.items(): residual_exprs.extend( - chain.from_iterable(self._build_residual(e, t) - for e in as_tuple(exprs) - )) - + chain.from_iterable( + self._build_residual(e, t) for e in as_tuple(exprs) + ) + ) self._F_exprs = tuple(sorted( residual_exprs, key=lambda e: not isinstance(e, EssentialBC) )) @@ -671,7 +684,9 @@ def _build_residual(self, expr, target): class InitialGuess: """ - Enforce initial guess to satisfy essential BCs. + Metadata passed to `LinearSolveExpr` to define the initial guess + symbolic expressions, enforcing the initial guess to satisfy essential + boundary conditions. """ def __init__(self, target, exprs, arrays): self.target = target @@ -680,13 +695,12 @@ def __init__(self, target, exprs, arrays): @property def exprs(self): - """ - """ return self._exprs def _build_exprs(self, exprs): """ - Return a list of initial guess expressions. + Return a list of initial guess symbolic expressions + that satisfy essential boundary conditions. """ self._exprs = tuple([ eq for eq in diff --git a/examples/petsc/cfd/01_navierstokes.py b/examples/petsc/cfd/01_navierstokes.py index 7797321045..7bd59a0c83 100644 --- a/examples/petsc/cfd/01_navierstokes.py +++ b/examples/petsc/cfd/01_navierstokes.py @@ -288,7 +288,7 @@ def neumann_right(eq, subdomain): tentu = PETScSolve([eq_u1]+bc_petsc_u1, u1.forward) tentv = PETScSolve([eq_v1]+bc_petsc_v1, v1.forward) -exprs = tentu + tentv + eqn_p + [update_u, update_v] + bc_u1 + bc_v1 +exprs = [tentu, tentv, eqn_p, update_u, update_v, bc_u1, bc_v1] with switchconfig(language='petsc'): op = Operator(exprs) diff --git a/examples/petsc/seismic/01_staggered_acoustic.py b/examples/petsc/seismic/01_staggered_acoustic.py index 4f0c91ed49..fc9e75938d 100644 --- a/examples/petsc/seismic/01_staggered_acoustic.py +++ b/examples/petsc/seismic/01_staggered_acoustic.py @@ -65,7 +65,7 @@ petsc_p_2 = PETScSolve(p_2, target=p2.forward, solver_parameters={'ksp_rtol': 1e-7}) with switchconfig(language='petsc'): - op_2 = Operator(petsc_v_x_2 + petsc_v_z_2 + petsc_p_2 + src_p_2, opt='noop') + op_2 = Operator([petsc_v_x_2, petsc_v_z_2, petsc_p_2, src_p_2], opt='noop') op_2(time=src.time_range.num-1, dt=dt) norm_p2 = norm(p2) @@ -90,7 +90,7 @@ petsc_p_4 = PETScSolve(p_4, target=p4.forward, solver_parameters={'ksp_rtol': 1e-7}) with switchconfig(language='petsc'): - op_4 = Operator(petsc_v_x_4 + petsc_v_z_4 + petsc_p_4 + src_p_4, opt='noop') + op_4 = Operator([petsc_v_x_4, petsc_v_z_4, petsc_p_4, src_p_4], opt='noop') op_4(time=src.time_range.num-1, dt=dt) norm_p4 = norm(p4) diff --git a/tests/test_petsc.py b/tests/test_petsc.py index c5f9c42b43..2352a83b91 100644 --- a/tests/test_petsc.py +++ b/tests/test_petsc.py @@ -150,7 +150,7 @@ def test_multiple_petsc_solves(): petsc2 = PETScSolve(eqn2, f2) with switchconfig(language='petsc'): - op = Operator(petsc1+petsc2, opt='noop') + op = Operator([petsc1, petsc2], opt='noop') callable_roots = [meta_call.root for meta_call in op._func_table.values()] @@ -320,7 +320,7 @@ def test_petsc_struct(): eqn2 = Eq(f1, g1*mu2) with switchconfig(language='petsc'): - op = Operator([eqn2] + petsc1) + op = Operator([eqn2, petsc1]) arguments = op.arguments() @@ -507,7 +507,7 @@ def test_time_loop(): petsc5 = PETScSolve(eq5, v2) with switchconfig(language='petsc'): - op4 = Operator(petsc4 + petsc5) + op4 = Operator([petsc4, petsc5]) op4.apply(time_M=3) body4 = str(op4.body) @@ -609,6 +609,53 @@ def define(self, dimensions): assert np.allclose(u.data[-1, 1:-1], 4.0) # right +@skipif('petsc') +def test_jacobian(): + + class SubLeft(SubDomain): + name = 'subleft' + + def define(self, dimensions): + x, = dimensions + return {x: ('left', 1)} + + class SubRight(SubDomain): + name = 'subright' + + def define(self, dimensions): + x, = dimensions + return {x: ('right', 1)} + + sub1 = SubLeft() + sub2 = SubRight() + + grid = Grid(shape=(11,), subdomains=(sub1, sub2), dtype=np.float64) + + e = Function(name='e', grid=grid, space_order=2) + f = Function(name='f', grid=grid, space_order=2) + + bc_1 = EssentialBC(e, 1.0, subdomain=sub1) + bc_2 = EssentialBC(e, 2.0, subdomain=sub2) + + eq1 = Eq(e.laplace + e, f + 2.0) + + petsc = PETScSolve([eq1, bc_1, bc_2], target=e) + + jac = petsc.rhs.fielddata.jacobian + + assert jac.row_target == e + assert jac.col_target == e + + # 2 symbolic expressions for each each EssentialBC (One ZeroRow and one ZeroColumn). + # NOTE: this is likely to change when PetscSection + DMDA is supported + assert len(jac.matvecs) == 5 + # TODO: I think some internals are preventing symplification here? + assert str(jac.scdiag) == 'h_x*(1 - 2.0/h_x**2)' + + assert all(isinstance(m, EssentialBC) for m in jac.matvecs[:4]) + assert not isinstance(jac.matvecs[-1], EssentialBC) + + class TestCoupledLinear: # The coupled interface can be used even for uncoupled problems, meaning # the equations will be solved within a single matrix system. @@ -655,7 +702,7 @@ def test_coupled_vs_non_coupled(self, eq1, eq2, so): petsc2 = PETScSolve(eq2, target=g) with switchconfig(language='petsc'): - op1 = Operator(petsc1 + petsc2, opt='noop') + op1 = Operator([petsc1, petsc2], opt='noop') op1.apply() enorm1 = norm(e) @@ -689,9 +736,9 @@ def test_coupled_vs_non_coupled(self, eq1, eq2, so): assert len(callbacks2) == 6 # Check fielddata type - fielddata1 = petsc1[0].rhs.fielddata - fielddata2 = petsc2[0].rhs.fielddata - fielddata3 = petsc3[0].rhs.fielddata + fielddata1 = petsc1.rhs.fielddata + fielddata2 = petsc2.rhs.fielddata + fielddata3 = petsc3.rhs.fielddata assert isinstance(fielddata1, FieldData) assert isinstance(fielddata2, FieldData) @@ -818,7 +865,7 @@ def test_mixed_jacobian(self): petsc = PETScSolve({e: [eq1], g: [eq2]}) - jacobian = petsc[0].rhs.fielddata.jacobian + jacobian = petsc.rhs.fielddata.jacobian j00 = jacobian.get_submatrix(0, 0) j01 = jacobian.get_submatrix(0, 1) @@ -894,7 +941,7 @@ def test_coupling(self, eq1, eq2, j01_matvec, j10_matvec): petsc = PETScSolve({e: [eq1], g: [eq2]}) - jacobian = petsc[0].rhs.fielddata.jacobian + jacobian = petsc.rhs.fielddata.jacobian j01 = jacobian.get_submatrix(0, 1) j10 = jacobian.get_submatrix(1, 0) @@ -943,7 +990,7 @@ def test_jacobian_scaling_1D(self, eq1, eq2, so, scale): petsc = PETScSolve({e: [eq1], g: [eq2]}) - jacobian = petsc[0].rhs.fielddata.jacobian + jacobian = petsc.rhs.fielddata.jacobian j00 = jacobian.get_submatrix(0, 0) j11 = jacobian.get_submatrix(1, 1) @@ -993,7 +1040,7 @@ def test_jacobian_scaling_2D(self, eq1, eq2, so, scale): petsc = PETScSolve({e: [eq1], g: [eq2]}) - jacobian = petsc[0].rhs.fielddata.jacobian + jacobian = petsc.rhs.fielddata.jacobian j00 = jacobian.get_submatrix(0, 0) j11 = jacobian.get_submatrix(1, 1) @@ -1046,7 +1093,7 @@ def test_jacobian_scaling_3D(self, eq1, eq2, so, scale): petsc = PETScSolve({e: [eq1], g: [eq2]}) - jacobian = petsc[0].rhs.fielddata.jacobian + jacobian = petsc.rhs.fielddata.jacobian j00 = jacobian.get_submatrix(0, 0) j11 = jacobian.get_submatrix(1, 1) From 905def2ca9cbdfaaf7dfeb3e8e446e7b39349e4d Mon Sep 17 00:00:00 2001 From: ZoeLeibowitz Date: Fri, 13 Jun 2025 14:39:12 +0100 Subject: [PATCH 52/84] misc: Flake8 --- tests/test_petsc.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/tests/test_petsc.py b/tests/test_petsc.py index 2352a83b91..87e58c2fe0 100644 --- a/tests/test_petsc.py +++ b/tests/test_petsc.py @@ -625,12 +625,12 @@ class SubRight(SubDomain): def define(self, dimensions): x, = dimensions return {x: ('right', 1)} - + sub1 = SubLeft() sub2 = SubRight() grid = Grid(shape=(11,), subdomains=(sub1, sub2), dtype=np.float64) - + e = Function(name='e', grid=grid, space_order=2) f = Function(name='f', grid=grid, space_order=2) @@ -645,7 +645,7 @@ def define(self, dimensions): assert jac.row_target == e assert jac.col_target == e - + # 2 symbolic expressions for each each EssentialBC (One ZeroRow and one ZeroColumn). # NOTE: this is likely to change when PetscSection + DMDA is supported assert len(jac.matvecs) == 5 From 19fddbe7fbdfa680d5782f3261116c0d02d3d85f Mon Sep 17 00:00:00 2001 From: ZoeLeibowitz Date: Fri, 13 Jun 2025 14:54:14 +0100 Subject: [PATCH 53/84] misc: Fix exprs in 01_navierstokes.py due to change in PETScSolve output type --- examples/petsc/cfd/01_navierstokes.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/examples/petsc/cfd/01_navierstokes.py b/examples/petsc/cfd/01_navierstokes.py index 7bd59a0c83..1c678d977b 100644 --- a/examples/petsc/cfd/01_navierstokes.py +++ b/examples/petsc/cfd/01_navierstokes.py @@ -288,7 +288,7 @@ def neumann_right(eq, subdomain): tentu = PETScSolve([eq_u1]+bc_petsc_u1, u1.forward) tentv = PETScSolve([eq_v1]+bc_petsc_v1, v1.forward) -exprs = [tentu, tentv, eqn_p, update_u, update_v, bc_u1, bc_v1] +exprs = [tentu, tentv, eqn_p, update_u, update_v] + bc_u1 + bc_v1 with switchconfig(language='petsc'): op = Operator(exprs) From 8037cf9c002fe9d7db49b52d9667c4a533aa7c88 Mon Sep 17 00:00:00 2001 From: ZoeLeibowitz Date: Fri, 13 Jun 2025 16:04:27 +0100 Subject: [PATCH 54/84] misc: Clean up --- devito/logger.py | 4 ++-- devito/operator/operator.py | 6 +----- devito/petsc/iet/passes.py | 19 +------------------ devito/petsc/iet/routines.py | 2 -- devito/petsc/initialize.py | 7 +++---- examples/petsc/petsc_test.py | 2 +- 6 files changed, 8 insertions(+), 32 deletions(-) diff --git a/devito/logger.py b/devito/logger.py index c9834dbfe7..e5df4bb565 100644 --- a/devito/logger.py +++ b/devito/logger.py @@ -77,9 +77,9 @@ def set_log_level(level, comm=None): """ from devito import configuration from devito.mpi.distributed import MPI - + if comm is not None and configuration['mpi']: - if comm!= MPI.COMM_NULL and comm.rank != 0: + if comm != MPI.COMM_NULL and comm.rank != 0: logger.removeHandler(stream_handler) logger.addHandler(logging.NullHandler()) else: diff --git a/devito/operator/operator.py b/devito/operator/operator.py index 5216bde481..3a36952a58 100644 --- a/devito/operator/operator.py +++ b/devito/operator/operator.py @@ -38,7 +38,7 @@ from devito.petsc.iet.passes import lower_petsc from devito.petsc.clusters import petsc_preprocess -__all__ = ['Operator', 'SpecialOp'] +__all__ = ['Operator'] class Operator(Callable): @@ -1456,7 +1456,3 @@ def parse_kwargs(**kwargs): kwargs['subs'] = {k: sympify(v) for k, v in kwargs.get('subs', {}).items()} return kwargs - - -class SpecialOp(Operator): - pass diff --git a/devito/petsc/iet/passes.py b/devito/petsc/iet/passes.py index f77879cf10..bf9a2c12fe 100644 --- a/devito/petsc/iet/passes.py +++ b/devito/petsc/iet/passes.py @@ -93,17 +93,6 @@ def initialize(iet): help_string = c.Line(r'static char help[] = "This is help text.\n";') init_body = petsc_call('PetscInitialize', [Byref(argc), Byref(argv), Null, Help]) - - # print_comm_size = - - # size = c.Line('PetscMPIInt size;') - # get_size = c.Line('PetscCallMPI(MPI_Comm_size(comm,&(size)));') - # rank = c.Line('PetscMPIInt rank;') - # get_rank = c.Line('PetscCallMPI(MPI_Comm_rank(comm,&(rank)));') - # print_comm_size = c.Line('PetscSynchronizedPrintf(comm, "MPI_Comm_size: %d\\n", size);') - # print_comm_rank = c.Line('PetscSynchronizedPrintf(comm, "MPI_Comm_rank: %d\\n", rank);') - - init_body = CallableBody( body=(petsc_func_begin_user, help_string, init_body), retstmt=(Call('PetscFunctionReturn', arguments=[0]),) @@ -128,20 +117,14 @@ def make_core_petsc_calls(objs, grid, **kwargs): comm = 'PETSC_COMM_WORLD' call_mpi = petsc_call_mpi('MPI_Comm_size', [comm, Byref(objs['size'])]) - # from IPython import embed; embed() - # get_size = c.Line('PetscCallMPI(MPI_Comm_size(comm,&(size)));') + rank = c.Line('PetscMPIInt rank;') get_rank = c.Line('PetscCallMPI(MPI_Comm_rank(comm,&(rank)));') print_comm_size = c.Line('PetscSynchronizedPrintf(comm, "MPI_Comm_size: %d\\n", size);') - print_comm_rank = c.Line('PetscSynchronizedPrintf(comm, "MPI_Comm_rank: %d\\n", rank);') - - flush = c.Line('PetscSynchronizedFlush(comm, PETSC_STDOUT);') - return call_mpi, rank, get_rank, print_comm_size, flush, print_comm_rank, flush, BlankLine - # return call_mpi, BlankLine class Builder: diff --git a/devito/petsc/iet/routines.py b/devito/petsc/iet/routines.py index f0e6a26756..7f8651e6f3 100644 --- a/devito/petsc/iet/routines.py +++ b/devito/petsc/iet/routines.py @@ -1087,7 +1087,6 @@ def __init__(self, **kwargs): self.objs = kwargs.get('objs') self.sregistry = kwargs.get('sregistry') self.grid = kwargs.get('grid') - # from IPython import embed; embed() self.devito_mpi = kwargs['options'].get('mpi', False) self.fielddata = self.injectsolve.expr.rhs.fielddata self.solver_objs = self._build() @@ -1129,7 +1128,6 @@ def _build(self): 'callbackdm': CallbackDM(sreg.make_name(prefix='dm')), } # TODO: Devito MPI + PETSc testing - # from IPython import embed; embed() if self.devito_mpi: base_dict['comm'] = self.grid.distributor._obj_comm else: diff --git a/devito/petsc/initialize.py b/devito/petsc/initialize.py index 3fc419c26a..80e3c7520c 100644 --- a/devito/petsc/initialize.py +++ b/devito/petsc/initialize.py @@ -3,8 +3,7 @@ from ctypes import POINTER, cast, c_char import atexit -from devito import switchconfig -from devito.operator.operator import SpecialOp +from devito import Operator, switchconfig from devito.types import Symbol from devito.types.equation import PetscEq from devito.petsc.types import Initialize, Finalize @@ -22,11 +21,11 @@ def PetscInitialize(): # This would prevent circular imports when initializing during import # from the PETSc module. with switchconfig(language='petsc'): - op_init = SpecialOp( + op_init = Operator( [PetscEq(dummy, Initialize(dummy))], name='kernel_init', opt='noop' ) - op_finalize = SpecialOp( + op_finalize = Operator( [PetscEq(dummy, Finalize(dummy))], name='kernel_finalize', opt='noop' ) diff --git a/examples/petsc/petsc_test.py b/examples/petsc/petsc_test.py index eb74c639db..5d93669d5f 100644 --- a/examples/petsc/petsc_test.py +++ b/examples/petsc/petsc_test.py @@ -28,4 +28,4 @@ op = Operator(petsc) op.apply() -# print(op.ccode) +print(op.ccode) From 8c16f2e112743f896613f88b53d521eeb922d02b Mon Sep 17 00:00:00 2001 From: ZoeLeibowitz Date: Fri, 13 Jun 2025 16:51:56 +0100 Subject: [PATCH 55/84] misc: Utilise zero_vector function --- devito/petsc/iet/routines.py | 8 ++------ 1 file changed, 2 insertions(+), 6 deletions(-) diff --git a/devito/petsc/iet/routines.py b/devito/petsc/iet/routines.py index f87eaab78b..fe9cce66f5 100644 --- a/devito/petsc/iet/routines.py +++ b/devito/petsc/iet/routines.py @@ -188,9 +188,7 @@ def _create_matvec_body(self, body, jacobian): 'DMGetLocalVector', [dmda, Byref(ylocal)] ) - zero_ylocal_memory = petsc_call( - 'VecSet', [ylocal, 0.0] - ) + zero_ylocal_memory = self.zero_vector(ylocal) vec_get_array_y = petsc_call( 'VecGetArray', [ylocal, Byref(y_matvec._C_symbol)] @@ -730,9 +728,7 @@ def _whole_matvec_body(self): nonzero_submats = self.jacobian.nonzero_submatrices - zero_y_memory = petsc_call( - 'VecSet', [objs['Y'], 0.0] - ) + zero_y_memory = self.zero_vector(objs['Y']) calls = () for sm in nonzero_submats: From f3d4d2deacd95184d26b6ee0ed3ef21cb7136856 Mon Sep 17 00:00:00 2001 From: ZoeLeibowitz Date: Fri, 13 Jun 2025 22:56:59 +0100 Subject: [PATCH 56/84] misc: Move vecset to function inside iet/utils.py --- devito/petsc/iet/routines.py | 38 +++++---------------------- devito/petsc/iet/utils.py | 4 +++ devito/petsc/types/types.py | 9 +++++-- examples/petsc/cfd/01_navierstokes.py | 5 ++++ 4 files changed, 22 insertions(+), 34 deletions(-) diff --git a/devito/petsc/iet/routines.py b/devito/petsc/iet/routines.py index f87eaab78b..9686c91d5e 100644 --- a/devito/petsc/iet/routines.py +++ b/devito/petsc/iet/routines.py @@ -16,7 +16,7 @@ from devito.petsc.types import PETScArray, PetscBundle from devito.petsc.iet.nodes import (PETScCallable, FormFunctionCallback, MatShellSetOp, PetscMetaData) -from devito.petsc.iet.utils import petsc_call, petsc_struct +from devito.petsc.iet.utils import petsc_call, petsc_struct, zero_vector from devito.petsc.utils import solver_mapper from devito.petsc.types import (DM, Mat, CallbackVec, Vec, KSP, PC, SNES, PetscInt, StartPtr, PointerIS, PointerDM, VecScatter, @@ -99,13 +99,6 @@ def initialguesses(self): def user_struct_callback(self): return self._user_struct_callback - @property - def zero_memory(self): - """Indicates whether the memory of the output - vector should be set to zero before the computation - in the callback.""" - return True - @property def fielddata(self): return self.injectsolve.expr.rhs.fielddata @@ -169,7 +162,7 @@ def _create_matvec_body(self, body, jacobian): 'DMGetApplicationContext', [dmda, Byref(ctx._C_symbol)] ) - zero_y_memory = self.zero_vector(objs['Y']) + zero_y_memory = zero_vector(objs['Y']) if jacobian.zero_memory else None dm_get_local_xvec = petsc_call( 'DMGetLocalVector', [dmda, Byref(xlocal)] @@ -188,9 +181,7 @@ def _create_matvec_body(self, body, jacobian): 'DMGetLocalVector', [dmda, Byref(ylocal)] ) - zero_ylocal_memory = petsc_call( - 'VecSet', [ylocal, 0.0] - ) + zero_ylocal_memory = zero_vector(ylocal) vec_get_array_y = petsc_call( 'VecGetArray', [ylocal, Byref(y_matvec._C_symbol)] @@ -320,7 +311,7 @@ def _create_formfunc_body(self, body): 'DMGetApplicationContext', [dmda, Byref(ctx._C_symbol)] ) - zero_f_memory = self.zero_vector(objs['F']) + zero_f_memory = zero_vector(objs['F']) dm_get_local_xvec = petsc_call( 'DMGetLocalVector', [dmda, Byref(objs['xloc'])] @@ -652,12 +643,6 @@ def _uxreplace_efuncs(self): mapper.update({k: visitor.visit(v)}) return mapper - def zero_vector(self, vec): - """ - Zeros the memory of the output vector before computation - """ - return petsc_call('VecSet', [vec, 0.0]) if self.zero_memory else None - def dereference_funcs(self, struct, fields): return tuple( [Dereference(i, struct) for i in @@ -691,13 +676,6 @@ def main_matvec_callback(self): def main_formfunc_callback(self): return self._main_formfunc_callback - @property - def zero_memory(self): - """Indicates whether the memory of the output - vector should be set to zero before the computation - in the callback.""" - return False - def _make_core(self): for sm in self.fielddata.jacobian.nonzero_submatrices: self._make_matvec(sm, prefix=f'{sm.name}_MatMult') @@ -730,9 +708,7 @@ def _whole_matvec_body(self): nonzero_submats = self.jacobian.nonzero_submatrices - zero_y_memory = petsc_call( - 'VecSet', [objs['Y'], 0.0] - ) + zero_y_memory = zero_vector(objs['Y']) calls = () for sm in nonzero_submats: @@ -815,9 +791,7 @@ def _whole_formfunc_body(self, body): 'DMGetApplicationContext', [dmda, Byref(ctx._C_symbol)] ) - zero_f_memory = petsc_call( - 'VecSet', [objs['F'], 0.0] - ) + zero_f_memory = zero_vector(objs['F']) dm_get_local_xvec = petsc_call( 'DMGetLocalVector', [dmda, Byref(objs['xloc'])] diff --git a/devito/petsc/iet/utils.py b/devito/petsc/iet/utils.py index 99da0468ad..d143bcc8c9 100644 --- a/devito/petsc/iet/utils.py +++ b/devito/petsc/iet/utils.py @@ -18,6 +18,10 @@ def petsc_struct(name, fields, pname, liveness='lazy', modifier=None): modifier=modifier) +def zero_vector(vec): + return petsc_call('VecSet', [vec, 0.0]) + + # Mapping special Eq operations to their corresponding IET Expression subclass types. # These operations correspond to subclasses of Eq utilised within PETScSolve. petsc_iet_mapper = {OpPetsc: PetscMetaData} diff --git a/devito/petsc/types/types.py b/devito/petsc/types/types.py index b97b6ed406..4f235bbbde 100644 --- a/devito/petsc/types/types.py +++ b/devito/petsc/types/types.py @@ -382,6 +382,10 @@ def row_target(self): def col_target(self): return self.target + @property + def zero_memory(self): + return True + def _build_matvecs(self): matvecs = [] for eq in self.exprs: @@ -516,8 +520,9 @@ def __init__(self, name, matvecs, scdiag, row_target, self.col_idx = col_idx self.linear_idx = linear_idx - def is_diag(self): - return self.row_idx == self.col_idx + @property + def zero_memory(self): + return False def __repr__(self): return (f"") diff --git a/examples/petsc/cfd/01_navierstokes.py b/examples/petsc/cfd/01_navierstokes.py index 1c678d977b..bd591c1eee 100644 --- a/examples/petsc/cfd/01_navierstokes.py +++ b/examples/petsc/cfd/01_navierstokes.py @@ -297,3 +297,8 @@ def neumann_right(eq, subdomain): # Pressure norm check tol = 1e-3 assert np.sum((pn1.data[0]-pn1.data[1])**2/np.maximum(pn1.data[0]**2, 1e-10)) < tol + +from devito import norm +print(norm(u1)) +print(norm(v1)) +print(norm(pn1)) From 8d10b8718e6702023e84e3d6c3de8e5d59001218 Mon Sep 17 00:00:00 2001 From: ZoeLeibowitz Date: Fri, 13 Jun 2025 22:59:21 +0100 Subject: [PATCH 57/84] misc: Flake8 --- examples/petsc/cfd/01_navierstokes.py | 5 ----- 1 file changed, 5 deletions(-) diff --git a/examples/petsc/cfd/01_navierstokes.py b/examples/petsc/cfd/01_navierstokes.py index bd591c1eee..1c678d977b 100644 --- a/examples/petsc/cfd/01_navierstokes.py +++ b/examples/petsc/cfd/01_navierstokes.py @@ -297,8 +297,3 @@ def neumann_right(eq, subdomain): # Pressure norm check tol = 1e-3 assert np.sum((pn1.data[0]-pn1.data[1])**2/np.maximum(pn1.data[0]**2, 1e-10)) < tol - -from devito import norm -print(norm(u1)) -print(norm(v1)) -print(norm(pn1)) From 53b58a92e501b5e1db53a343187a84da445bc0f4 Mon Sep 17 00:00:00 2001 From: ZoeLeibowitz Date: Fri, 13 Jun 2025 23:50:10 +0100 Subject: [PATCH 58/84] misc: Clean up --- devito/petsc/iet/routines.py | 47 +++++++++++------------------------- devito/petsc/iet/utils.py | 45 ++++++++++++++++++++++++++++++++++ 2 files changed, 59 insertions(+), 33 deletions(-) diff --git a/devito/petsc/iet/routines.py b/devito/petsc/iet/routines.py index 9686c91d5e..ec78054094 100644 --- a/devito/petsc/iet/routines.py +++ b/devito/petsc/iet/routines.py @@ -16,7 +16,8 @@ from devito.petsc.types import PETScArray, PetscBundle from devito.petsc.iet.nodes import (PETScCallable, FormFunctionCallback, MatShellSetOp, PetscMetaData) -from devito.petsc.iet.utils import petsc_call, petsc_struct, zero_vector +from devito.petsc.iet.utils import (petsc_call, petsc_struct, zero_vector, + dereference_funcs, residual_bundle) from devito.petsc.utils import solver_mapper from devito.petsc.types import (DM, Mat, CallbackVec, Vec, KSP, PC, SNES, PetscInt, StartPtr, PointerIS, PointerDM, VecScatter, @@ -120,7 +121,7 @@ def _make_core(self): self._make_user_struct_callback() def _make_matvec(self, jacobian, prefix='MatMult'): - # Compile matvec `eqns` into an IET via recursive compilation + # Compile `matvecs` into an IET via recursive compilation matvecs = jacobian.matvecs irs, _ = self.rcompile( matvecs, options={'mpi': False}, sregistry=self.sregistry, @@ -251,7 +252,7 @@ def _create_matvec_body(self, body, jacobian): ) # Dereference function data in struct - derefs = self.dereference_funcs(ctx, fields) + derefs = dereference_funcs(ctx, fields) body = CallableBody( List(body=body), @@ -390,7 +391,7 @@ def _create_formfunc_body(self, body): ) # Dereference function data in struct - derefs = self.dereference_funcs(ctx, fields) + derefs = dereference_funcs(ctx, fields) body = CallableBody( List(body=body), @@ -500,7 +501,7 @@ def _create_form_rhs_body(self, body): ) # Dereference function data in struct - derefs = self.dereference_funcs(ctx, fields) + derefs = dereference_funcs(ctx, fields) body = CallableBody( List(body=[body]), @@ -578,7 +579,7 @@ def _create_initial_guess_body(self, body): ) # Dereference function data in struct - derefs = self.dereference_funcs(ctx, fields) + derefs = dereference_funcs(ctx, fields) body = CallableBody( List(body=[body]), @@ -643,12 +644,6 @@ def _uxreplace_efuncs(self): mapper.update({k: visitor.visit(v)}) return mapper - def dereference_funcs(self, struct, fields): - return tuple( - [Dereference(i, struct) for i in - fields if isinstance(i.function, AbstractFunction)] - ) - class CCBBuilder(CBBuilder): def __init__(self, **kwargs): @@ -749,17 +744,17 @@ def _whole_matvec_body(self): def _make_whole_formfunc(self): F_exprs = self.fielddata.residual.F_exprs - # Compile formfunc `eqns` into an IET via recursive compilation - irs_formfunc, _ = self.rcompile( + # Compile `F_exprs` into an IET via recursive compilation + irs, _ = self.rcompile( F_exprs, options={'mpi': False}, sregistry=self.sregistry, concretize_mapper=self.concretize_mapper ) - body_formfunc = self._whole_formfunc_body(List(body=irs_formfunc.uiet.body)) + body = self._whole_formfunc_body(List(body=irs.uiet.body)) objs = self.objs cb = PETScCallable( self.sregistry.make_name(prefix='WholeFormFunc'), - body_formfunc, + body, retval=objs['err'], parameters=(objs['snes'], objs['X'], objs['F'], objs['dummyptr']) ) @@ -783,7 +778,8 @@ def _whole_formfunc_body(self, body): bundles = sobjs['bundles'] fbundle = bundles['f'] xbundle = bundles['x'] - body = self.residual_bundle(body, bundles) + + body = residual_bundle(body, bundles) dm_cast = DummyExpr(dmda, DMCast(objs['dummyptr']), init=True) @@ -870,7 +866,7 @@ def _whole_formfunc_body(self, body): ) # Dereference function data in struct - derefs = self.dereference_funcs(ctx, fields) + derefs = dereference_funcs(ctx, fields) f_soa = PointerCast(fbundle) x_soa = PointerCast(xbundle) @@ -1034,21 +1030,6 @@ def _submat_callback_body(self): retstmt=(Call('PetscFunctionReturn', arguments=[0]),) ) - def residual_bundle(self, body, bundles): - mapper = bundles['bundle_mapper'] - indexeds = FindSymbols('indexeds').visit(body) - subs = {} - - for i in indexeds: - if i.base in mapper: - bundle = mapper[i.base] - index = bundles['target_indices'][i.function.target] - index = (index,) + i.indices - subs[i] = bundle.__getitem__(index) - - body = Uxreplace(subs).visit(body) - return body - class BaseObjectBuilder: """ diff --git a/devito/petsc/iet/utils.py b/devito/petsc/iet/utils.py index d143bcc8c9..3402cfe28f 100644 --- a/devito/petsc/iet/utils.py +++ b/devito/petsc/iet/utils.py @@ -1,5 +1,7 @@ from devito.petsc.iet.nodes import PetscMetaData, PETScCall from devito.ir.equations import OpPetsc +from devito.ir.iet import Dereference, FindSymbols, Uxreplace +from devito.types.basic import AbstractFunction def petsc_call(specific_call, call_args): @@ -19,9 +21,52 @@ def petsc_struct(name, fields, pname, liveness='lazy', modifier=None): def zero_vector(vec): + """ + Set all entries of a PETSc vector to zero. + """ return petsc_call('VecSet', [vec, 0.0]) +def dereference_funcs(struct, fields): + """ + Dereference AbstractFunctions from a struct. + """ + return tuple( + [Dereference(i, struct) for i in + fields if isinstance(i.function, AbstractFunction)] + ) + + +def residual_bundle(body, bundles): + """ + Replaces PetscArrays in `body` with PetscBundle struct field accesses + (e.g., f_v[ix][iy] -> f_bundle[ix][iy].v). + + Example: + f_v[ix][iy] = x_v[ix][iy]; + f_u[ix][iy] = x_u[ix][iy]; + becomes: + f_bundle[ix][iy].v = x_bundle[ix][iy].v; + f_bundle[ix][iy].u = x_bundle[ix][iy].u; + + NOTE: This is used because the data is interleaved for + multi-component DMDAs in PETSc. + """ + mapper = bundles['bundle_mapper'] + indexeds = FindSymbols('indexeds').visit(body) + subs = {} + + for i in indexeds: + if i.base in mapper: + bundle = mapper[i.base] + index = bundles['target_indices'][i.function.target] + index = (index,) + i.indices + subs[i] = bundle.__getitem__(index) + + body = Uxreplace(subs).visit(body) + return 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 = {OpPetsc: PetscMetaData} From 857eb9cebc711a36b3608e56433ec78f78e46e04 Mon Sep 17 00:00:00 2001 From: ZoeLeibowitz Date: Sun, 15 Jun 2025 20:55:46 +0100 Subject: [PATCH 59/84] misc: Add more tests --- tests/test_petsc.py | 83 +++++++++++++++++++++++++++++++++++++++++++-- 1 file changed, 81 insertions(+), 2 deletions(-) diff --git a/tests/test_petsc.py b/tests/test_petsc.py index 87e58c2fe0..678b1071c8 100644 --- a/tests/test_petsc.py +++ b/tests/test_petsc.py @@ -12,7 +12,8 @@ from devito.passes.iet.languages.C import CDataManager from devito.petsc.types import (DM, Mat, Vec, PetscMPIInt, KSP, PC, KSPConvergedReason, PETScArray, - LinearSolveExpr, FieldData, MultipleFieldData) + LinearSolveExpr, FieldData, MultipleFieldData, + SubMatrixBlock) from devito.petsc.solve import PETScSolve, EssentialBC from devito.petsc.iet.nodes import Expression from devito.petsc.initialize import PetscInitialize @@ -647,7 +648,7 @@ def define(self, dimensions): assert jac.col_target == e # 2 symbolic expressions for each each EssentialBC (One ZeroRow and one ZeroColumn). - # NOTE: this is likely to change when PetscSection + DMDA is supported + # NOTE: This is likely to change when PetscSection + DMDA is supported assert len(jac.matvecs) == 5 # TODO: I think some internals are preventing symplification here? assert str(jac.scdiag) == 'h_x*(1 - 2.0/h_x**2)' @@ -656,6 +657,48 @@ def define(self, dimensions): assert not isinstance(jac.matvecs[-1], EssentialBC) +@skipif('petsc') +def test_residual(): + class SubLeft(SubDomain): + name = 'subleft' + + def define(self, dimensions): + x, = dimensions + return {x: ('left', 1)} + + class SubRight(SubDomain): + name = 'subright' + + def define(self, dimensions): + x, = dimensions + return {x: ('right', 1)} + + sub1 = SubLeft() + sub2 = SubRight() + + grid = Grid(shape=(11,), subdomains=(sub1, sub2), dtype=np.float64) + + e = Function(name='e', grid=grid, space_order=2) + f = Function(name='f', grid=grid, space_order=2) + + bc_1 = EssentialBC(e, 1.0, subdomain=sub1) + bc_2 = EssentialBC(e, 2.0, subdomain=sub2) + + eq1 = Eq(e.laplace + e, f + 2.0) + + petsc = PETScSolve([eq1, bc_1, bc_2], target=e) + + res = petsc.rhs.fielddata.residual + + assert res.target == e + # NOTE: This is likely to change when PetscSection + DMDA is supported + assert len(res.F_exprs) == 5 + assert len(res.b_exprs) == 3 + + assert not res.time_mapper + assert str(res.scdiag) == 'h_x*(1 - 2.0/h_x**2)' + + class TestCoupledLinear: # The coupled interface can be used even for uncoupled problems, meaning # the equations will be solved within a single matrix system. @@ -872,6 +915,42 @@ def test_mixed_jacobian(self): j10 = jacobian.get_submatrix(1, 0) j11 = jacobian.get_submatrix(1, 1) + # Check type of each submatrix is a SubMatrixBlock + assert isinstance(j00, SubMatrixBlock) + assert isinstance(j01, SubMatrixBlock) + assert isinstance(j10, SubMatrixBlock) + assert isinstance(j11, SubMatrixBlock) + + assert j00.name == 'J00' + assert j01.name == 'J01' + assert j10.name == 'J10' + assert j11.name == 'J11' + + assert j00.row_target == e + assert j01.row_target == e + assert j10.row_target == g + assert j11.row_target == g + + assert j00.col_target == e + assert j01.col_target == g + assert j10.col_target == e + assert j11.col_target == g + + assert j00.row_idx == 0 + assert j01.row_idx == 0 + assert j10.row_idx == 1 + assert j11.row_idx == 1 + + assert j00.col_idx == 0 + assert j01.col_idx == 1 + assert j10.col_idx == 0 + assert j11.col_idx == 1 + + assert j00.linear_idx == 0 + assert j01.linear_idx == 1 + assert j10.linear_idx == 2 + assert j11.linear_idx == 3 + # Check the number of submatrices assert jacobian.n_submatrices == 4 From 7a5b10fd84275d21b8056ece5bb1af46cb0303c5 Mon Sep 17 00:00:00 2001 From: ZoeLeibowitz Date: Mon, 16 Jun 2025 09:44:06 +0100 Subject: [PATCH 60/84] mpi: Start parallel tests --- tests/test_petsc.py | 12 ++++++++++++ 1 file changed, 12 insertions(+) diff --git a/tests/test_petsc.py b/tests/test_petsc.py index 87e58c2fe0..daa6b6d734 100644 --- a/tests/test_petsc.py +++ b/tests/test_petsc.py @@ -1232,3 +1232,15 @@ def define(self, dimensions): + 'MATOP_MULT,(void (*)(void))J00_MatMult0)' in str(create) # TODO: Test mixed, time dependent solvers + + +class TestMPI: + + @pytest.mark.parallel(mode=4) + def test_laplacian(self, mode): + """ + """ + grid = Grid(shape=(4,)) + + f = Function(name='f') + From 467c19babb6e28484ea11e75c955a8d77e7174a8 Mon Sep 17 00:00:00 2001 From: ZoeLeibowitz Date: Mon, 16 Jun 2025 10:12:26 +0100 Subject: [PATCH 61/84] add to petsc laplacian test --- devito/petsc/iet/passes.py | 3 +- tests/test_petsc.py | 80 ++++++++++++++++++++++++++++++++++++-- 2 files changed, 78 insertions(+), 5 deletions(-) diff --git a/devito/petsc/iet/passes.py b/devito/petsc/iet/passes.py index bf9a2c12fe..f6af3f0a4f 100644 --- a/devito/petsc/iet/passes.py +++ b/devito/petsc/iet/passes.py @@ -124,7 +124,8 @@ def make_core_petsc_calls(objs, grid, **kwargs): print_comm_rank = c.Line('PetscSynchronizedPrintf(comm, "MPI_Comm_rank: %d\\n", rank);') flush = c.Line('PetscSynchronizedFlush(comm, PETSC_STDOUT);') - return call_mpi, rank, get_rank, print_comm_size, flush, print_comm_rank, flush, BlankLine + # return call_mpi, rank, get_rank, print_comm_size, flush, print_comm_rank, flush, BlankLine + return call_mpi, BlankLine class Builder: diff --git a/tests/test_petsc.py b/tests/test_petsc.py index daa6b6d734..222de35282 100644 --- a/tests/test_petsc.py +++ b/tests/test_petsc.py @@ -1236,11 +1236,83 @@ def define(self, dimensions): class TestMPI: - @pytest.mark.parallel(mode=4) - def test_laplacian(self, mode): + # @pytest.mark.parallel(mode=4) + @skipif('petsc') + def test_laplacian(self): """ """ - grid = Grid(shape=(4,)) - f = Function(name='f') + # Subdomains to implement BCs + class SubLeft(SubDomain): + name = 'subleft' + + def define(self, dimensions): + x, = dimensions + return {x: ('left', 1)} + + + class SubRight(SubDomain): + name = 'subright' + + def define(self, dimensions): + x, = dimensions + return {x: ('right', 1)} + + + sub1 = SubLeft() + sub2 = SubRight() + subdomains = (sub1, sub2,) + + + def exact(x): + return -np.float64(np.exp(x)) + + + # n = 9, 17, 33, 65, 129, 257 + n_values = [2**k + 1 for k in range(3, 9)] + n_values = [9] + dx = np.array([1./(n-1) for n in n_values]) + errors = [] + + for n in n_values: + + grid = Grid( + shape=(n,), subdomains=subdomains, dtype=np.float64 + ) + + u = Function(name='u', grid=grid, space_order=2) + f = Function(name='f', grid=grid, space_order=2) + bc = Function(name='bc', grid=grid, space_order=2) + + eqn = Eq(-u.laplace, f, subdomain=grid.interior) + + X = np.linspace(0, 1.0, n).astype(np.float64) + f.data[:] = np.float64(np.exp(X)) + + bc.data[0] = -np.float64(1.0) # u(0) = -1 + bc.data[-1] = -np.float64(np.exp(1.0)) # u(1) = -e + + # Create boundary condition expressions using subdomains + bcs = [EssentialBC(u, bc, subdomain=sub1)] + bcs += [EssentialBC(u, bc, subdomain=sub2)] + + exprs = [eqn] + bcs + petsc = PETScSolve(exprs, target=u, solver_parameters={'ksp_rtol': 1e-10}) + + op = Operator(petsc, language='petsc') + op.apply() + + # u_exact = Function(name='u_exact', grid=grid, space_order=2) + # u_exact.data[:] = exact(X) + + # diff = u_exact.data[:] - u.data[:] + # u_diff_norm = np.linalg.norm(diff, ord=np.inf) + # u_error = u_diff_norm / np.linalg.norm(u_exact.data[:], ord=np.inf) + # errors.append(u_error) + + # Expected norms computed "manually" from sequential runs + norm_u = norm(u) + assert norm_u == 2.0 + # Expected norms computed "manually" from sequential runs + # assert np.isclose(norm(ux), 7003.098, rtol=1.e-4) \ No newline at end of file From d2e3eb59260fecc9020160db495a3e6ee7a2d110 Mon Sep 17 00:00:00 2001 From: ZoeLeibowitz Date: Mon, 16 Jun 2025 10:28:53 +0100 Subject: [PATCH 62/84] mpi tests --- tests/test_petsc.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/tests/test_petsc.py b/tests/test_petsc.py index b01dcac2f1..c0e36919dc 100644 --- a/tests/test_petsc.py +++ b/tests/test_petsc.py @@ -1317,7 +1317,8 @@ class TestMPI: # @pytest.mark.parallel(mode=4) @skipif('petsc') - def test_laplacian(self): + @pytest.mark.parallel(mode=[(2)]) + def test_laplacian(self, mode): """ """ @@ -1391,7 +1392,7 @@ def exact(x): # Expected norms computed "manually" from sequential runs norm_u = norm(u) - assert norm_u == 2.0 + assert norm_u == 5.467052700706644 # Expected norms computed "manually" from sequential runs # assert np.isclose(norm(ux), 7003.098, rtol=1.e-4) \ No newline at end of file From 6165373e003ee4e17c01a0b7e66c61fd07122f53 Mon Sep 17 00:00:00 2001 From: ZoeLeibowitz Date: Mon, 16 Jun 2025 12:48:13 +0100 Subject: [PATCH 63/84] edit test --- examples/petsc/petsc_test.py | 2 +- tests/test_petsc.py | 71 +++++++++++------------------------- 2 files changed, 23 insertions(+), 50 deletions(-) diff --git a/examples/petsc/petsc_test.py b/examples/petsc/petsc_test.py index 5d93669d5f..eb74c639db 100644 --- a/examples/petsc/petsc_test.py +++ b/examples/petsc/petsc_test.py @@ -28,4 +28,4 @@ op = Operator(petsc) op.apply() -print(op.ccode) +# print(op.ccode) diff --git a/tests/test_petsc.py b/tests/test_petsc.py index c0e36919dc..72823c4985 100644 --- a/tests/test_petsc.py +++ b/tests/test_petsc.py @@ -1314,15 +1314,13 @@ def define(self, dimensions): class TestMPI: + # TODO: Add test for DMDACreate() in parallel - # @pytest.mark.parallel(mode=4) @skipif('petsc') - @pytest.mark.parallel(mode=[(2)]) + @pytest.mark.parallel(mode=1) def test_laplacian(self, mode): """ """ - - # Subdomains to implement BCs class SubLeft(SubDomain): name = 'subleft' @@ -1330,7 +1328,6 @@ def define(self, dimensions): x, = dimensions return {x: ('left', 1)} - class SubRight(SubDomain): name = 'subright' @@ -1343,56 +1340,32 @@ def define(self, dimensions): sub2 = SubRight() subdomains = (sub1, sub2,) + n = 9 - def exact(x): - return -np.float64(np.exp(x)) - - - # n = 9, 17, 33, 65, 129, 257 - n_values = [2**k + 1 for k in range(3, 9)] - n_values = [9] - dx = np.array([1./(n-1) for n in n_values]) - errors = [] - - for n in n_values: - - grid = Grid( - shape=(n,), subdomains=subdomains, dtype=np.float64 - ) + grid = Grid( + shape=(n,), subdomains=subdomains, dtype=np.float64 + ) - u = Function(name='u', grid=grid, space_order=2) - f = Function(name='f', grid=grid, space_order=2) - bc = Function(name='bc', grid=grid, space_order=2) - - eqn = Eq(-u.laplace, f, subdomain=grid.interior) - - X = np.linspace(0, 1.0, n).astype(np.float64) - f.data[:] = np.float64(np.exp(X)) - - bc.data[0] = -np.float64(1.0) # u(0) = -1 - bc.data[-1] = -np.float64(np.exp(1.0)) # u(1) = -e + u = Function(name='u', grid=grid, space_order=2) + f = Function(name='f', grid=grid, space_order=2) + bc = Function(name='bc', grid=grid, space_order=2) - # Create boundary condition expressions using subdomains - bcs = [EssentialBC(u, bc, subdomain=sub1)] - bcs += [EssentialBC(u, bc, subdomain=sub2)] + eqn = Eq(-u.laplace, f, subdomain=grid.interior) - exprs = [eqn] + bcs - petsc = PETScSolve(exprs, target=u, solver_parameters={'ksp_rtol': 1e-10}) + X = np.linspace(0, 1.0, n).astype(np.float64) + f.data[:] = np.float64(np.exp(X)) - op = Operator(petsc, language='petsc') - op.apply() + bc.data[0] = -np.float64(1.0) # u(0) = -1 + bc.data[-1] = -np.float64(np.exp(1.0)) # u(1) = -e - # u_exact = Function(name='u_exact', grid=grid, space_order=2) - # u_exact.data[:] = exact(X) + # Create boundary condition expressions using subdomains + bcs = [EssentialBC(u, bc, subdomain=sub1)] + bcs += [EssentialBC(u, bc, subdomain=sub2)] - # diff = u_exact.data[:] - u.data[:] - # u_diff_norm = np.linalg.norm(diff, ord=np.inf) - # u_error = u_diff_norm / np.linalg.norm(u_exact.data[:], ord=np.inf) - # errors.append(u_error) + petsc = PETScSolve([eqn] + bcs, target=u, solver_parameters={'ksp_rtol': 1e-10}) - # Expected norms computed "manually" from sequential runs - norm_u = norm(u) - assert norm_u == 5.467052700706644 + op = Operator(petsc, language='petsc') + op.apply() - # Expected norms computed "manually" from sequential runs - # assert np.isclose(norm(ux), 7003.098, rtol=1.e-4) \ No newline at end of file + # Expected norm computed "manually" from sequential run + assert np.isclose(norm(u), 5.467052700706644, rtol=1e-15, atol=1e-15) From 0b43ea68b5c9d6fc31d6b726abc078af729a8e2b Mon Sep 17 00:00:00 2001 From: ZoeLeibowitz Date: Mon, 16 Jun 2025 18:33:55 +0100 Subject: [PATCH 64/84] tests: Add petsc parallel test --- requirements-testing.txt | 1 + tests/test_petsc.py | 33 +++++++++++++++++++++++++-------- 2 files changed, 26 insertions(+), 8 deletions(-) diff --git a/requirements-testing.txt b/requirements-testing.txt index d10923cb00..2acf0e4f11 100644 --- a/requirements-testing.txt +++ b/requirements-testing.txt @@ -1,6 +1,7 @@ pytest>=7.2,<8.5 pytest-runner<6.0.2 pytest-cov<6.1.2 +pytest-order flake8-pyproject>=1.2.3,<1.2.4 nbval<0.11.1 scipy<1.15.3 diff --git a/tests/test_petsc.py b/tests/test_petsc.py index 72823c4985..1132ebd5fb 100644 --- a/tests/test_petsc.py +++ b/tests/test_petsc.py @@ -20,6 +20,7 @@ @skipif('petsc') +@pytest.mark.order(0) def test_petsc_initialization(): # TODO: Temporary workaround until PETSc is automatically # initialized @@ -28,6 +29,14 @@ def test_petsc_initialization(): PetscInitialize() +@skipif('petsc') +@pytest.mark.parallel(mode=[1, 2, 4, 6]) +def test_petsc_initialization_parallel(mode): + configuration['compiler'] = 'custom' + os.environ['CC'] = 'mpicc' + PetscInitialize() + + @skipif('petsc') def test_petsc_local_object(): """ @@ -1316,11 +1325,22 @@ def define(self, dimensions): class TestMPI: # TODO: Add test for DMDACreate() in parallel + @pytest.mark.parametrize('nx, unorm', [ + (17, 7.441506654790017), + (33, 10.317652759863675), + (65, 14.445123374862874), + (129, 20.32492895656658), + (257, 28.67050632840985) + ]) @skipif('petsc') - @pytest.mark.parallel(mode=1) - def test_laplacian(self, mode): + @pytest.mark.parallel(mode=[1, 2, 4, 8]) + def test_laplacian_1d(self, nx, unorm, mode): """ """ + configuration['compiler'] = 'custom' + os.environ['CC'] = 'mpicc' + PetscInitialize() + class SubLeft(SubDomain): name = 'subleft' @@ -1335,15 +1355,12 @@ def define(self, dimensions): x, = dimensions return {x: ('right', 1)} - sub1 = SubLeft() sub2 = SubRight() subdomains = (sub1, sub2,) - n = 9 - grid = Grid( - shape=(n,), subdomains=subdomains, dtype=np.float64 + shape=(nx,), subdomains=subdomains, dtype=np.float64 ) u = Function(name='u', grid=grid, space_order=2) @@ -1352,7 +1369,7 @@ def define(self, dimensions): eqn = Eq(-u.laplace, f, subdomain=grid.interior) - X = np.linspace(0, 1.0, n).astype(np.float64) + X = np.linspace(0, 1.0, nx).astype(np.float64) f.data[:] = np.float64(np.exp(X)) bc.data[0] = -np.float64(1.0) # u(0) = -1 @@ -1368,4 +1385,4 @@ def define(self, dimensions): op.apply() # Expected norm computed "manually" from sequential run - assert np.isclose(norm(u), 5.467052700706644, rtol=1e-15, atol=1e-15) + assert np.isclose(norm(u), unorm, rtol=1e-13, atol=1e-13) From 62d004538fc638792c5f3bd4f1e59e842264768a Mon Sep 17 00:00:00 2001 From: ZoeLeibowitz Date: Mon, 16 Jun 2025 18:43:09 +0100 Subject: [PATCH 65/84] clean up --- devito/petsc/iet/passes.py | 10 +--------- 1 file changed, 1 insertion(+), 9 deletions(-) diff --git a/devito/petsc/iet/passes.py b/devito/petsc/iet/passes.py index f6af3f0a4f..f358696c7a 100644 --- a/devito/petsc/iet/passes.py +++ b/devito/petsc/iet/passes.py @@ -115,16 +115,8 @@ def make_core_petsc_calls(objs, grid, **kwargs): comm = grid.distributor._obj_comm else: comm = 'PETSC_COMM_WORLD' - - call_mpi = petsc_call_mpi('MPI_Comm_size', [comm, Byref(objs['size'])]) - - rank = c.Line('PetscMPIInt rank;') - get_rank = c.Line('PetscCallMPI(MPI_Comm_rank(comm,&(rank)));') - print_comm_size = c.Line('PetscSynchronizedPrintf(comm, "MPI_Comm_size: %d\\n", size);') - print_comm_rank = c.Line('PetscSynchronizedPrintf(comm, "MPI_Comm_rank: %d\\n", rank);') - flush = c.Line('PetscSynchronizedFlush(comm, PETSC_STDOUT);') - # return call_mpi, rank, get_rank, print_comm_size, flush, print_comm_rank, flush, BlankLine + call_mpi = petsc_call_mpi('MPI_Comm_size', [comm, Byref(objs['size'])]) return call_mpi, BlankLine From bd963794aa1879e50825af431c58fc490312873b Mon Sep 17 00:00:00 2001 From: ZoeLeibowitz Date: Mon, 16 Jun 2025 19:42:05 +0100 Subject: [PATCH 66/84] misc: Simplify comm extraction in lower_petsc --- devito/petsc/iet/passes.py | 18 ++++++++---------- devito/petsc/iet/routines.py | 12 ++++++------ examples/petsc/petsc_test.py | 2 -- 3 files changed, 14 insertions(+), 18 deletions(-) diff --git a/devito/petsc/iet/passes.py b/devito/petsc/iet/passes.py index f358696c7a..f20f30991f 100644 --- a/devito/petsc/iet/passes.py +++ b/devito/petsc/iet/passes.py @@ -53,9 +53,11 @@ def lower_petsc(iet, **kwargs): if len(unique_grids) > 1: raise ValueError("All PETScSolves must use the same Grid, but multiple found.") grid = unique_grids.pop() + devito_mpi = kwargs['options'].get('mpi', False) + comm = grid.distributor._obj_comm if devito_mpi else 'PETSC_COMM_WORLD' # Create core PETSc calls (not specific to each PETScSolve) - core = make_core_petsc_calls(objs, grid, **kwargs) + core = make_core_petsc_calls(objs, comm) setup = [] subs = {} @@ -63,7 +65,7 @@ def lower_petsc(iet, **kwargs): for iters, (injectsolve,) in injectsolve_mapper.items(): - builder = Builder(injectsolve, objs, iters, grid, **kwargs) + builder = Builder(injectsolve, objs, iters, grid, comm, **kwargs) setup.extend(builder.solversetup.calls) @@ -109,13 +111,7 @@ def finalize(iet): return iet._rebuild(body=finalize_body) -def make_core_petsc_calls(objs, grid, **kwargs): - devito_mpi = kwargs['options'].get('mpi', False) - if devito_mpi: - comm = grid.distributor._obj_comm - else: - comm = 'PETSC_COMM_WORLD' - +def make_core_petsc_calls(objs, comm): call_mpi = petsc_call_mpi('MPI_Comm_size', [comm, Byref(objs['size'])]) return call_mpi, BlankLine @@ -129,11 +125,12 @@ class Builder: returning subclasses of the objects initialised in __init__, depending on the properties of `injectsolve`. """ - def __init__(self, injectsolve, objs, iters, grid, **kwargs): + def __init__(self, injectsolve, objs, iters, grid, comm, **kwargs): self.injectsolve = injectsolve self.objs = objs self.iters = iters self.grid = grid + self.comm = comm self.kwargs = kwargs self.coupled = isinstance(injectsolve.expr.rhs.fielddata, MultipleFieldData) self.args = { @@ -141,6 +138,7 @@ def __init__(self, injectsolve, objs, iters, grid, **kwargs): 'objs': self.objs, 'iters': self.iters, 'grid': self.grid, + 'comm': self.comm, **self.kwargs } self.args['solver_objs'] = self.objbuilder.solver_objs diff --git a/devito/petsc/iet/routines.py b/devito/petsc/iet/routines.py index 3ecf52462d..f999b13e71 100644 --- a/devito/petsc/iet/routines.py +++ b/devito/petsc/iet/routines.py @@ -1041,8 +1041,7 @@ def __init__(self, **kwargs): self.injectsolve = kwargs.get('injectsolve') self.objs = kwargs.get('objs') self.sregistry = kwargs.get('sregistry') - self.grid = kwargs.get('grid') - self.devito_mpi = kwargs['options'].get('mpi', False) + self.comm = kwargs.get('comm') self.fielddata = self.injectsolve.expr.rhs.fielddata self.solver_objs = self._build() @@ -1083,10 +1082,11 @@ def _build(self): 'callbackdm': CallbackDM(sreg.make_name(prefix='dm')), } # TODO: Devito MPI + PETSc testing - if self.devito_mpi: - base_dict['comm'] = self.grid.distributor._obj_comm - else: - base_dict['comm'] = 'PETSC_COMM_WORLD' + # if self.devito_mpi: + # base_dict['comm'] = self.grid.distributor._obj_comm + # else: + # base_dict['comm'] = 'PETSC_COMM_WORLD' + base_dict['comm'] = self.comm self._target_dependent(base_dict) return self._extend_build(base_dict) diff --git a/examples/petsc/petsc_test.py b/examples/petsc/petsc_test.py index eb74c639db..821a3569fc 100644 --- a/examples/petsc/petsc_test.py +++ b/examples/petsc/petsc_test.py @@ -27,5 +27,3 @@ with switchconfig(language='petsc'): op = Operator(petsc) op.apply() - -# print(op.ccode) From ce5cf82672d006bc07fabfc563a17010fb25a1bf Mon Sep 17 00:00:00 2001 From: ZoeLeibowitz Date: Mon, 16 Jun 2025 19:57:09 +0100 Subject: [PATCH 67/84] clean up --- devito/petsc/iet/passes.py | 6 ++---- devito/petsc/iet/routines.py | 5 ----- examples/petsc/petsc_test.py | 2 ++ tests/test_petsc.py | 1 + 4 files changed, 5 insertions(+), 9 deletions(-) diff --git a/devito/petsc/iet/passes.py b/devito/petsc/iet/passes.py index f20f30991f..b4859dda16 100644 --- a/devito/petsc/iet/passes.py +++ b/devito/petsc/iet/passes.py @@ -65,7 +65,7 @@ def lower_petsc(iet, **kwargs): for iters, (injectsolve,) in injectsolve_mapper.items(): - builder = Builder(injectsolve, objs, iters, grid, comm, **kwargs) + builder = Builder(injectsolve, objs, iters, comm, **kwargs) setup.extend(builder.solversetup.calls) @@ -125,11 +125,10 @@ class Builder: returning subclasses of the objects initialised in __init__, depending on the properties of `injectsolve`. """ - def __init__(self, injectsolve, objs, iters, grid, comm, **kwargs): + def __init__(self, injectsolve, objs, iters, comm, **kwargs): self.injectsolve = injectsolve self.objs = objs self.iters = iters - self.grid = grid self.comm = comm self.kwargs = kwargs self.coupled = isinstance(injectsolve.expr.rhs.fielddata, MultipleFieldData) @@ -137,7 +136,6 @@ def __init__(self, injectsolve, objs, iters, grid, comm, **kwargs): 'injectsolve': self.injectsolve, 'objs': self.objs, 'iters': self.iters, - 'grid': self.grid, 'comm': self.comm, **self.kwargs } diff --git a/devito/petsc/iet/routines.py b/devito/petsc/iet/routines.py index f999b13e71..da6ded1f9a 100644 --- a/devito/petsc/iet/routines.py +++ b/devito/petsc/iet/routines.py @@ -1081,11 +1081,6 @@ def _build(self): 'dmda': DM(sreg.make_name(prefix='da'), dofs=len(targets)), 'callbackdm': CallbackDM(sreg.make_name(prefix='dm')), } - # TODO: Devito MPI + PETSc testing - # if self.devito_mpi: - # base_dict['comm'] = self.grid.distributor._obj_comm - # else: - # base_dict['comm'] = 'PETSC_COMM_WORLD' base_dict['comm'] = self.comm self._target_dependent(base_dict) return self._extend_build(base_dict) diff --git a/examples/petsc/petsc_test.py b/examples/petsc/petsc_test.py index 821a3569fc..5d93669d5f 100644 --- a/examples/petsc/petsc_test.py +++ b/examples/petsc/petsc_test.py @@ -27,3 +27,5 @@ with switchconfig(language='petsc'): op = Operator(petsc) op.apply() + +print(op.ccode) diff --git a/tests/test_petsc.py b/tests/test_petsc.py index 1132ebd5fb..a2f9d8b8ec 100644 --- a/tests/test_petsc.py +++ b/tests/test_petsc.py @@ -1385,4 +1385,5 @@ def define(self, dimensions): op.apply() # Expected norm computed "manually" from sequential run + # What rtol and atol should be used? assert np.isclose(norm(u), unorm, rtol=1e-13, atol=1e-13) From b7c4082554e77ce10321d0532c2b75706da5fa54 Mon Sep 17 00:00:00 2001 From: ZoeLeibowitz Date: Mon, 16 Jun 2025 21:23:23 +0100 Subject: [PATCH 68/84] trigger petsc CI --- .github/workflows/pytest-petsc.yml | 2 ++ 1 file changed, 2 insertions(+) diff --git a/.github/workflows/pytest-petsc.yml b/.github/workflows/pytest-petsc.yml index 2bf55b5049..502e9b2139 100644 --- a/.github/workflows/pytest-petsc.yml +++ b/.github/workflows/pytest-petsc.yml @@ -11,10 +11,12 @@ on: branches: - main - petsc + - biharmonic pull_request: branches: - main - petsc + - biharmonic jobs: pytest: From e485c9c002b0f13c0002636e52f5a1f64ad05198 Mon Sep 17 00:00:00 2001 From: ZoeLeibowitz Date: Tue, 17 Jun 2025 12:47:57 +0100 Subject: [PATCH 69/84] misc: Add todo: --- devito/petsc/iet/routines.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/devito/petsc/iet/routines.py b/devito/petsc/iet/routines.py index da6ded1f9a..9249c7a557 100644 --- a/devito/petsc/iet/routines.py +++ b/devito/petsc/iet/routines.py @@ -1389,7 +1389,10 @@ def _create_dmda(self, dmda): # Number of degrees of freedom per node args.append(dmda.dofs) # "Stencil width" -> size of overlap + # TODO: Instead, this probably should be + # extracted from fielddata.target._size_outhalo? stencil_width = self.fielddata.space_order + args.append(stencil_width) args.extend([objs['Null']]*nspace_dims) From df4e638f738d3c3047ee7d3728447358f7bae363 Mon Sep 17 00:00:00 2001 From: ZoeLeibowitz Date: Tue, 17 Jun 2025 15:56:51 +0100 Subject: [PATCH 70/84] address some of ed's comments --- requirements-testing.txt | 1 - tests/test_petsc.py | 42 +++++++++++++++++++++++++--------------- 2 files changed, 26 insertions(+), 17 deletions(-) diff --git a/requirements-testing.txt b/requirements-testing.txt index 2acf0e4f11..d10923cb00 100644 --- a/requirements-testing.txt +++ b/requirements-testing.txt @@ -1,7 +1,6 @@ pytest>=7.2,<8.5 pytest-runner<6.0.2 pytest-cov<6.1.2 -pytest-order flake8-pyproject>=1.2.3,<1.2.4 nbval<0.11.1 scipy<1.15.3 diff --git a/tests/test_petsc.py b/tests/test_petsc.py index a2f9d8b8ec..371901be76 100644 --- a/tests/test_petsc.py +++ b/tests/test_petsc.py @@ -20,7 +20,7 @@ @skipif('petsc') -@pytest.mark.order(0) +@pytest.fixture(scope='session', autouse=True) def test_petsc_initialization(): # TODO: Temporary workaround until PETSc is automatically # initialized @@ -1341,27 +1341,37 @@ def test_laplacian_1d(self, nx, unorm, mode): os.environ['CC'] = 'mpicc' PetscInitialize() - class SubLeft(SubDomain): - name = 'subleft' + # class SubLeft(SubDomain): + # name = 'subleft' - def define(self, dimensions): - x, = dimensions - return {x: ('left', 1)} + # def define(self, dimensions): + # x, = dimensions + # return {x: ('left', 1)} + + # class SubRight(SubDomain): + # name = 'subright' - class SubRight(SubDomain): - name = 'subright' + # def define(self, dimensions): + # x, = dimensions + # return {x: ('right', 1)} + + # grid = Grid(shape=(nx,), dtype=np.float64) + + # sub1 = SubLeft(grid=grid) + # sub2 = SubRight(grid=grid) + + class SubSide(SubDomain): + def __init__(self, side='left', grid=None): + self.side = side + self.name = f'sub{side}' + super().__init__(grid=grid) def define(self, dimensions): x, = dimensions - return {x: ('right', 1)} - - sub1 = SubLeft() - sub2 = SubRight() - subdomains = (sub1, sub2,) + return {x: (self.side, 1)} - grid = Grid( - shape=(nx,), subdomains=subdomains, dtype=np.float64 - ) + grid = Grid(shape=(nx,), dtype=np.float64) + sub1, sub2 = [SubSide(side=s, grid=grid) for s in ('left', 'right')] u = Function(name='u', grid=grid, space_order=2) f = Function(name='f', grid=grid, space_order=2) From 72e8222769caaf84382f6aecb92f41e2c1c75d91 Mon Sep 17 00:00:00 2001 From: ZoeLeibowitz Date: Tue, 17 Jun 2025 17:07:23 +0100 Subject: [PATCH 71/84] misc: Address comments --- tests/test_petsc.py | 32 ++++++-------------------------- 1 file changed, 6 insertions(+), 26 deletions(-) diff --git a/tests/test_petsc.py b/tests/test_petsc.py index 371901be76..3f64a79c82 100644 --- a/tests/test_petsc.py +++ b/tests/test_petsc.py @@ -1333,7 +1333,7 @@ class TestMPI: (257, 28.67050632840985) ]) @skipif('petsc') - @pytest.mark.parallel(mode=[1, 2, 4, 8]) + @pytest.mark.parallel(mode=[2, 4, 8]) def test_laplacian_1d(self, nx, unorm, mode): """ """ @@ -1341,25 +1341,6 @@ def test_laplacian_1d(self, nx, unorm, mode): os.environ['CC'] = 'mpicc' PetscInitialize() - # class SubLeft(SubDomain): - # name = 'subleft' - - # def define(self, dimensions): - # x, = dimensions - # return {x: ('left', 1)} - - # class SubRight(SubDomain): - # name = 'subright' - - # def define(self, dimensions): - # x, = dimensions - # return {x: ('right', 1)} - - # grid = Grid(shape=(nx,), dtype=np.float64) - - # sub1 = SubLeft(grid=grid) - # sub2 = SubRight(grid=grid) - class SubSide(SubDomain): def __init__(self, side='left', grid=None): self.side = side @@ -1375,19 +1356,18 @@ def define(self, dimensions): u = Function(name='u', grid=grid, space_order=2) f = Function(name='f', grid=grid, space_order=2) - bc = Function(name='bc', grid=grid, space_order=2) + + u0 = Constant(name='u0', value=-1.0, dtype=np.float64) + u1 = Constant(name='u1', value=-np.exp(1.0), dtype=np.float64) eqn = Eq(-u.laplace, f, subdomain=grid.interior) X = np.linspace(0, 1.0, nx).astype(np.float64) f.data[:] = np.float64(np.exp(X)) - bc.data[0] = -np.float64(1.0) # u(0) = -1 - bc.data[-1] = -np.float64(np.exp(1.0)) # u(1) = -e - # Create boundary condition expressions using subdomains - bcs = [EssentialBC(u, bc, subdomain=sub1)] - bcs += [EssentialBC(u, bc, subdomain=sub2)] + bcs = [EssentialBC(u, u0, subdomain=sub1)] + bcs += [EssentialBC(u, u1, subdomain=sub2)] petsc = PETScSolve([eqn] + bcs, target=u, solver_parameters={'ksp_rtol': 1e-10}) From 16be19a9a4b0372bf6078007013a36bdb7137fc2 Mon Sep 17 00:00:00 2001 From: ZoeLeibowitz Date: Fri, 20 Jun 2025 12:28:51 +0100 Subject: [PATCH 72/84] dsl/compiler: Add PETSc logging instrastructure and tests --- devito/operator/operator.py | 6 +- devito/operator/profiling.py | 36 +++++- devito/petsc/iet/logging.py | 77 ++++++++++++ devito/petsc/iet/passes.py | 51 +++++--- devito/petsc/iet/routines.py | 42 ++++--- devito/petsc/logging.py | 180 +++++++++++++++++++++++++++++ devito/petsc/solve.py | 16 ++- devito/petsc/types/object.py | 2 + devito/petsc/types/types.py | 10 +- devito/symbolics/extended_sympy.py | 6 +- tests/test_petsc.py | 108 ++++++++++++++++- 11 files changed, 483 insertions(+), 51 deletions(-) create mode 100644 devito/petsc/iet/logging.py create mode 100644 devito/petsc/logging.py diff --git a/devito/operator/operator.py b/devito/operator/operator.py index 3a36952a58..61464120c5 100644 --- a/devito/operator/operator.py +++ b/devito/operator/operator.py @@ -194,7 +194,7 @@ def _sanitize_exprs(cls, expressions, **kwargs): @classmethod def _build(cls, expressions, **kwargs): # Python- (i.e., compile-) and C-level (i.e., run-time) performance - profiler = create_profile('timers') + profiler = create_profile('timers', kwargs['language']) # Lower the input expressions into an IET irs, byproduct = cls._lower(expressions, profiler=profiler, **kwargs) @@ -1004,7 +1004,9 @@ def _emit_apply_profiling(self, args): elapsed = fround(self._profiler.py_timers['apply']) info(f"Operator `{self.name}` ran in {elapsed:.2f} s") - summary = self._profiler.summary(args, self._dtype, reduce_over=elapsed) + summary = self._profiler.summary( + args, self._dtype, self.parameters, reduce_over=elapsed + ) if not is_log_enabled_for('PERF'): # Do not waste time diff --git a/devito/operator/profiling.py b/devito/operator/profiling.py index ae2f50e171..360d506d17 100644 --- a/devito/operator/profiling.py +++ b/devito/operator/profiling.py @@ -17,6 +17,7 @@ from devito.parameters import configuration from devito.symbolics import subs_op_args from devito.tools import DefaultOrderedDict, flatten +from devito.petsc.logging import PetscSummary __all__ = ['create_profile'] @@ -42,7 +43,7 @@ class Profiler: _attempted_init = False - def __init__(self, name): + def __init__(self, name, language): self.name = name # Operation reductions observed in sections @@ -55,6 +56,9 @@ def __init__(self, name): # Python-level timers self.py_timers = OrderedDict() + # For language specific summaries + self.language = language + self._attempted_init = True def analyze(self, iet): @@ -179,7 +183,7 @@ def record_ops_variation(self, initial, final): def all_sections(self): return list(self._sections) + flatten(self._subsections.values()) - def summary(self, args, dtype, reduce_over=None): + def summary(self, args, dtype, params, reduce_over=None): """ Return a PerformanceSummary of the profiled sections. @@ -277,7 +281,7 @@ def _allgather_from_comm(self, comm, time, ops, points, traffic, sops, itershape return list(zip(times, opss, pointss, traffics, sops, itershapess)) # Override basic summary so that arguments other than runtime are computed. - def summary(self, args, dtype, reduce_over=None): + def summary(self, args, dtype, params, reduce_over=None): grid = args.grid comm = args.comm @@ -338,6 +342,11 @@ def summary(self, args, dtype, reduce_over=None): # data transfers) summary.add_glb_fdlike('fdlike-nosetup', points, reduce_over_nosetup) + # Add the language specific summary + summary.add_language_summary( + self.language, + language_summary_mapper[self.language](params) + ) return summary @@ -478,6 +487,16 @@ def add_glb_fdlike(self, key, points, time): self.globals[key] = PerfEntry(time, None, gpointss, None, None, None) + def add_language_summary(self, lang, summary): + """ + Register a language specific summary (e.g., PetscSummary) + and dynamically add a property to access it via perf_summary.. + """ + # TODO: Consider renaming `PerformanceSummary` to something more generic + # (e.g., `Summary`), or separating `PetscSummary` entirely from + # `PerformanceSummary`. + setattr(self, lang, summary) + @property def globals_all(self): v0 = self.globals['vanilla'] @@ -503,7 +522,7 @@ def timings(self): return OrderedDict([(k, v.time) for k, v in self.items()]) -def create_profile(name): +def create_profile(name, language): """Create a new Profiler.""" if configuration['log-level'] in ['DEBUG', 'PERF'] and \ configuration['profiling'] == 'basic': @@ -511,13 +530,13 @@ def create_profile(name): level = 'advanced' else: level = configuration['profiling'] - profiler = profiler_registry[level](name) + profiler = profiler_registry[level](name, language) if profiler._attempted_init: return profiler else: warning(f"Couldn't set up `{level}` profiler; reverting to 'advanced'") - profiler = profiler_registry['advanced'](name) + profiler = profiler_registry['advanced'](name, language) # We expect the `advanced` profiler to always initialize successfully assert profiler._attempted_init return profiler @@ -533,3 +552,8 @@ def create_profile(name): 'advisor': AdvisorProfiler } """Profiling levels.""" + + +language_summary_mapper = { + 'petsc': PetscSummary +} diff --git a/devito/petsc/iet/logging.py b/devito/petsc/iet/logging.py new file mode 100644 index 0000000000..6a1c74d767 --- /dev/null +++ b/devito/petsc/iet/logging.py @@ -0,0 +1,77 @@ +from functools import cached_property + +from devito.symbolics import Byref, FieldFromPointer +from devito.ir.iet import DummyExpr +from devito.logger import PERF + +from devito.petsc.iet.utils import petsc_call +from devito.petsc.logging import petsc_return_variable_dict, PetscInfo + + +class PetscLogger: + """ + Class for PETSc loggers that collect solver related statistics. + """ + def __init__(self, level, **kwargs): + self.sobjs = kwargs.get('solver_objs') + self.sreg = kwargs.get('sregistry') + self.section_mapper = kwargs.get('section_mapper', {}) + self.injectsolve = kwargs.get('injectsolve', None) + + self.function_list = [] + + if level <= PERF: + self.function_list.extend([ + 'kspgetiterationnumber', + 'snesgetiterationnumber' + ]) + + # TODO: To be extended with if level <= DEBUG: ... + + name = self.sreg.make_name(prefix='petscinfo') + pname = self.sreg.make_name(prefix='petscprofiler') + + self.statstruct = PetscInfo( + name, pname, self.logobjs, self.sobjs, + self.section_mapper, self.injectsolve, + self.function_list + ) + + @cached_property + def logobjs(self): + """ + Create PETSc objects specifically needed for logging solver statistics. + """ + return { + info.name: info.variable_type( + self.sreg.make_name(prefix=info.output_param) + ) + for func_name in self.function_list + for info in [petsc_return_variable_dict[func_name]] + } + + @cached_property + def calls(self): + """ + Generate the PETSc calls that will be injected into the C code to + extract solver statistics. + """ + struct = self.statstruct + calls = [] + for param in self.function_list: + param = petsc_return_variable_dict[param] + + inputs = [] + for i in param.input_params: + inputs.append(self.sobjs[i]) + + logobj = self.logobjs[param.name] + + calls.append( + petsc_call(param.name, inputs + [Byref(logobj)]) + ) + # TODO: Perform a PetscCIntCast here? + expr = DummyExpr(FieldFromPointer(logobj._C_symbol, struct), logobj._C_symbol) + calls.append(expr) + + return tuple(calls) diff --git a/devito/petsc/iet/passes.py b/devito/petsc/iet/passes.py index b4859dda16..8d7042bb6c 100644 --- a/devito/petsc/iet/passes.py +++ b/devito/petsc/iet/passes.py @@ -5,7 +5,7 @@ from devito.passes.iet.engine import iet_pass from devito.ir.iet import (Transformer, MapNodes, Iteration, BlankLine, DummyExpr, CallableBody, List, Call, Callable, - FindNodes) + FindNodes, Section) from devito.symbolics import Byref, Macro, FieldFromPointer from devito.types import Symbol, Scalar from devito.types.basic import DataSymbol @@ -22,8 +22,11 @@ CoupledObjectBuilder, BaseSetup, CoupledSetup, Solver, CoupledSolver, TimeDependent, NonTimeDependent) +from devito.petsc.iet.logging import PetscLogger from devito.petsc.iet.utils import petsc_call, petsc_call_mpi +import devito.logger as dl + @iet_pass def lower_petsc(iet, **kwargs): @@ -63,14 +66,17 @@ def lower_petsc(iet, **kwargs): subs = {} efuncs = {} + # Map PETScSolve to its Section (for logging) + section_mapper = MapNodes(Section, PetscMetaData, 'groupby').visit(iet) + for iters, (injectsolve,) in injectsolve_mapper.items(): - builder = Builder(injectsolve, objs, iters, comm, **kwargs) + builder = Builder(injectsolve, objs, iters, comm, section_mapper, **kwargs) setup.extend(builder.solversetup.calls) # Transform the spatial iteration loop with the calls to execute the solver - subs.update({builder.solve.spatial_body: builder.solve.calls}) + subs.update({builder.solve.spatial_body: builder.calls}) efuncs.update(builder.cbbuilder.efuncs) @@ -78,7 +84,7 @@ def lower_petsc(iet, **kwargs): iet = Transformer(subs).visit(iet) - body = core + tuple(setup) + (BlankLine,) + iet.body.body + body = core + tuple(setup) + iet.body.body body = iet.body._rebuild(body=body) iet = iet._rebuild(body=body) metadata = {**core_metadata(), 'efuncs': tuple(efuncs.values())} @@ -125,49 +131,64 @@ class Builder: returning subclasses of the objects initialised in __init__, depending on the properties of `injectsolve`. """ - def __init__(self, injectsolve, objs, iters, comm, **kwargs): + def __init__(self, injectsolve, objs, iters, comm, section_mapper, **kwargs): self.injectsolve = injectsolve self.objs = objs self.iters = iters self.comm = comm + self.section_mapper = section_mapper self.kwargs = kwargs self.coupled = isinstance(injectsolve.expr.rhs.fielddata, MultipleFieldData) - self.args = { + self.common_kwargs = { 'injectsolve': self.injectsolve, 'objs': self.objs, 'iters': self.iters, 'comm': self.comm, + 'section_mapper': self.section_mapper, **self.kwargs } - self.args['solver_objs'] = self.objbuilder.solver_objs - self.args['timedep'] = self.timedep - self.args['cbbuilder'] = self.cbbuilder + self.common_kwargs['solver_objs'] = self.objbuilder.solver_objs + self.common_kwargs['timedep'] = self.timedep + self.common_kwargs['cbbuilder'] = self.cbbuilder + self.common_kwargs['logger'] = self.logger @cached_property def objbuilder(self): return ( - CoupledObjectBuilder(**self.args) + CoupledObjectBuilder(**self.common_kwargs) if self.coupled else - BaseObjectBuilder(**self.args) + BaseObjectBuilder(**self.common_kwargs) ) @cached_property def timedep(self): time_mapper = self.injectsolve.expr.rhs.time_mapper timedep_class = TimeDependent if time_mapper else NonTimeDependent - return timedep_class(**self.args) + return timedep_class(**self.common_kwargs) @cached_property def cbbuilder(self): - return CCBBuilder(**self.args) if self.coupled else CBBuilder(**self.args) + return CCBBuilder(**self.common_kwargs) \ + if self.coupled else CBBuilder(**self.common_kwargs) @cached_property def solversetup(self): - return CoupledSetup(**self.args) if self.coupled else BaseSetup(**self.args) + return CoupledSetup(**self.common_kwargs) \ + if self.coupled else BaseSetup(**self.common_kwargs) @cached_property def solve(self): - return CoupledSolver(**self.args) if self.coupled else Solver(**self.args) + return CoupledSolver(**self.common_kwargs) \ + if self.coupled else Solver(**self.common_kwargs) + + @cached_property + def logger(self): + log_level = dl.logger.level + return PetscLogger(log_level, **self.common_kwargs) + + @cached_property + def calls(self): + return List(body=self.solve.calls+self.logger.calls) def populate_matrix_context(efuncs, objs): diff --git a/devito/petsc/iet/routines.py b/devito/petsc/iet/routines.py index 9249c7a557..d21af4be6b 100644 --- a/devito/petsc/iet/routines.py +++ b/devito/petsc/iet/routines.py @@ -7,22 +7,21 @@ retrieve_iteration_tree, filter_iterations, Iteration, PointerCast) from devito.symbolics import (Byref, FieldFromPointer, cast, VOID, - FieldFromComposite, IntDiv, Deref, Mod) + FieldFromComposite, IntDiv, Deref, Mod, String) from devito.symbolics.unevaluation import Mul from devito.types.basic import AbstractFunction from devito.types import Temp, Dimension from devito.tools import filter_ordered -from devito.petsc.types import PETScArray, PetscBundle from devito.petsc.iet.nodes import (PETScCallable, FormFunctionCallback, MatShellSetOp, PetscMetaData) from devito.petsc.iet.utils import (petsc_call, petsc_struct, zero_vector, dereference_funcs, residual_bundle) from devito.petsc.utils import solver_mapper -from devito.petsc.types import (DM, Mat, CallbackVec, Vec, KSP, PC, SNES, - PetscInt, StartPtr, PointerIS, PointerDM, VecScatter, - DMCast, JacobianStruct, - SubMatrixStruct, CallbackDM) +from devito.petsc.types import (PETScArray, PetscBundle, DM, Mat, CallbackVec, Vec, + KSP, PC, SNES, PetscInt, StartPtr, PointerIS, PointerDM, + VecScatter, DMCast, JacobianStruct, SubMatrixStruct, + CallbackDM) class CBBuilder: @@ -1068,6 +1067,10 @@ def _build(self): """ sreg = self.sregistry targets = self.fielddata.targets + + snes_name = sreg.make_name(prefix='snes') + options_prefix = self.injectsolve.expr.rhs.options_prefix + base_dict = { 'Jac': Mat(sreg.make_name(prefix='J')), 'xglobal': Vec(sreg.make_name(prefix='xglobal')), @@ -1076,10 +1079,12 @@ def _build(self): 'blocal': CallbackVec(sreg.make_name(prefix='blocal')), 'ksp': KSP(sreg.make_name(prefix='ksp')), 'pc': PC(sreg.make_name(prefix='pc')), - 'snes': SNES(sreg.make_name(prefix='snes')), + 'snes': SNES(snes_name), 'localsize': PetscInt(sreg.make_name(prefix='localsize')), 'dmda': DM(sreg.make_name(prefix='da'), dofs=len(targets)), 'callbackdm': CallbackDM(sreg.make_name(prefix='dm')), + 'snesprefix': String((options_prefix or '') + '_'), + 'options_prefix': options_prefix, } base_dict['comm'] = self.comm self._target_dependent(base_dict) @@ -1239,6 +1244,10 @@ def _setup(self): snes_create = petsc_call('SNESCreate', [sobjs['comm'], Byref(sobjs['snes'])]) + snes_options_prefix = petsc_call( + 'SNESSetOptionsPrefix', [sobjs['snes'], sobjs['snesprefix']] + ) if sobjs['options_prefix'] else None + snes_set_dm = petsc_call('SNESSetDM', [sobjs['snes'], dmda]) create_matrix = petsc_call('DMCreateMatrix', [dmda, Byref(sobjs['Jac'])]) @@ -1327,6 +1336,7 @@ def _setup(self): base_setup = dmda_calls + ( snes_create, + snes_options_prefix, snes_set_dm, create_matrix, snes_set_jac, @@ -1412,11 +1422,14 @@ def _setup(self): sobjs = self.solver_objs dmda = sobjs['dmda'] - solver_params = self.injectsolve.expr.rhs.solver_parameters snes_create = petsc_call('SNESCreate', [sobjs['comm'], Byref(sobjs['snes'])]) + snes_options_prefix = petsc_call( + 'SNESSetOptionsPrefix', [sobjs['snes'], sobjs['snesprefix']] + ) if sobjs['options_prefix'] else None + snes_set_dm = petsc_call('SNESSetDM', [sobjs['snes'], dmda]) create_matrix = petsc_call('DMCreateMatrix', [dmda, Byref(sobjs['Jac'])]) @@ -1530,6 +1543,7 @@ def _setup(self): coupled_setup = dmda_calls + ( snes_create, + snes_options_prefix, snes_set_dm, create_matrix, snes_set_jac, @@ -1618,7 +1632,7 @@ def _execute_solve(self): vec_reset_array, BlankLine, ) - return List(body=run_solver_calls) + return run_solver_calls @cached_property def spatial_body(self): @@ -1701,15 +1715,7 @@ def _execute_solve(self): snes_solve = (petsc_call('SNESSolve', [sobjs['snes'], objs['Null'], xglob]),) - return List( - body=( - (struct_assignment,) - + pre_solve - + snes_solve - + post_solve - + (BlankLine,) - ) - ) + return (struct_assignment,) + pre_solve + snes_solve + post_solve + (BlankLine,) class NonTimeDependent: diff --git a/devito/petsc/logging.py b/devito/petsc/logging.py new file mode 100644 index 0000000000..f552b69969 --- /dev/null +++ b/devito/petsc/logging.py @@ -0,0 +1,180 @@ +from collections import namedtuple, OrderedDict +from dataclasses import dataclass + +from devito.types import CompositeObject + +from devito.petsc.types import PetscInt +from devito.petsc.utils import petsc_type_mappings + + +class PetscEntry: + def __init__(self, **kwargs): + self.kwargs = kwargs + for k, v in self.kwargs.items(): + setattr(self, k, v) + self._properties = {k.lower(): v for k, v in kwargs.items()} + + def __getitem__(self, key): + return self._properties[key.lower()] + + def __repr__(self): + return f"PetscEntry({', '.join(f'{k}={v}' for k, v in self.kwargs.items())})" + + +class PetscSummary(dict): + """ + A summary of PETSc statistics collected for all solver runs + associated with a single operator during execution. + """ + PetscKey = namedtuple('PetscKey', 'name options_prefix') + + def __init__(self, params, *args, **kwargs): + super().__init__(*args, **kwargs) + + self.petscinfos = [i for i in params if isinstance(i, PetscInfo)] + + # Gather all unique PETSc function names across all PetscInfo objects + self._functions = list(dict.fromkeys( + petsc_return_variable_dict[key].name + for struct in self.petscinfos + for key in struct.function_list + )) + self._property_name_map = {} + # Dynamically create a property on this class for each PETSc function + self._add_properties() + + # Initialize the summary by adding PETSc information from each PetscInfo + # object (each corresponding to an individual PETScSolve) + for i in self.petscinfos: + self.add_info(i) + + def add_info(self, petscinfo): + """ + For a given PetscInfo object, create a key + and entry and add it to the PetscSummary. + """ + key = self.PetscKey(*petscinfo.summary_key) + entry = self.petsc_entry(petscinfo) + self[key] = entry + + def petsc_entry(self, petscinfo): + """ + Create a named tuple entry for the given PetscInfo object, + containing the values for each PETSc function call. + """ + funcs = self._functions + values = tuple(getattr(petscinfo, c) for c in funcs) + return PetscEntry(**{k: v for k, v in zip(funcs, values)}) + + def _add_properties(self): + """ + For each function name in `self._functions` (e.g., 'KSPGetIterationNumber'), + dynamically add a property to the class with the same name. + + Each property returns an OrderedDict that maps each PetscKey to the + result of looking up that function on the corresponding PetscEntry, + if the function exists on that entry. + """ + def make_property(function): + def getter(self): + return OrderedDict( + (k, getattr(v, function)) + for k, v in self.items() + # Only include entries that have the function + if hasattr(v, function) + ) + return property(getter) + + for f in self._functions: + # Inject the new property onto the class itself + setattr(self.__class__, f, make_property(f)) + self._property_name_map[f.lower()] = f + + def get_entry(self, name, options_prefix=None): + """ + Retrieve a single PetscEntry for a given name + and options_prefix. + """ + key = self.PetscKey(name, options_prefix) + if key not in self: + raise ValueError( + f"No PETSc information for:" + f" name='{name}'" + f" options_prefix='{options_prefix}'" + ) + return self[key] + + def __getitem__(self, key): + if isinstance(key, str): + # Allow case insensitive key access + original = self._property_name_map.get(key.lower()) + if original: + return getattr(self, original) + raise KeyError(f"No PETSc function named '{key}'") + elif isinstance(key, tuple) and len(key) == 2: + # Allow tuple keys (name, options_prefix) + key = self.PetscKey(*key) + return super().__getitem__(key) + + +class PetscInfo(CompositeObject): + + __rargs__ = ('name', 'pname', 'logobjs', 'sobjs', 'section_mapper', + 'injectsolve', 'function_list') + + def __init__(self, name, pname, logobjs, sobjs, section_mapper, + injectsolve, function_list): + + self.logobjs = logobjs + self.sobjs = sobjs + self.section_mapper = section_mapper + self.injectsolve = injectsolve + self.function_list = function_list + + mapper = {v: k for k, v in petsc_type_mappings.items()} + fields = [(str(i), mapper[str(i._C_ctype)]) for i in logobjs.values()] + super().__init__(name, pname, fields) + + @property + def section(self): + section = self.section_mapper.items() + return next((k[0].name for k, v in section if self.injectsolve in v), None) + + @property + def summary_key(self): + return (self.section, self.sobjs['options_prefix']) + + def __getattr__(self, attr): + if attr in self.logobjs.keys(): + return getattr(self.value._obj, self.logobjs[attr].name) + raise AttributeError(f"{attr} not found in PETSc return variables") + + +@dataclass +class PetscReturnVariable: + name: str + variable_type: None + input_params: list + output_param: str + + +# NOTE: +# In the future, this dictionary should be generated automatically from PETSc internals. +# For now, it serves as the reference for PETSc function metadata. +# If any of the PETSc function signatures change (e.g., names, input/output parameters), +# this dictionary must be updated accordingly. + +petsc_return_variable_dict = { + 'kspgetiterationnumber': PetscReturnVariable( + name='KSPGetIterationNumber', + variable_type=PetscInt, + input_params=['ksp'], + output_param='kspiter' + ), + 'snesgetiterationnumber': PetscReturnVariable( + name='SNESGetIterationNumber', + variable_type=PetscInt, + input_params=['snes'], + output_param='snesiter', + ) +} diff --git a/devito/petsc/solve.py b/devito/petsc/solve.py index 1f4423707e..55633d40fd 100644 --- a/devito/petsc/solve.py +++ b/devito/petsc/solve.py @@ -12,7 +12,7 @@ __all__ = ['PETScSolve'] -def PETScSolve(target_exprs, target=None, solver_parameters=None): +def PETScSolve(target_exprs, target=None, solver_parameters=None, options_prefix=None): """ Returns a symbolic expression representing a linear PETSc solver, enriched with all the necessary metadata for execution within an `Operator`. @@ -51,16 +51,21 @@ def PETScSolve(target_exprs, target=None, solver_parameters=None): This can be passed directly to a Devito Operator. """ if target is not None: - return InjectSolve(solver_parameters, {target: target_exprs}).build_expr() + return InjectSolve( + solver_parameters, {target: target_exprs}, options_prefix + ).build_expr() else: - return InjectMixedSolve(solver_parameters, target_exprs).build_expr() + return InjectMixedSolve( + solver_parameters, target_exprs, options_prefix + ).build_expr() class InjectSolve: - def __init__(self, solver_parameters=None, target_exprs=None): + def __init__(self, solver_parameters=None, target_exprs=None, options_prefix=None): self.solver_params = solver_parameters self.time_mapper = None self.target_exprs = target_exprs + self.options_prefix = options_prefix def build_expr(self): target, funcs, fielddata = self.linear_solve_args() @@ -71,7 +76,8 @@ def build_expr(self): self.solver_params, fielddata=fielddata, time_mapper=self.time_mapper, - localinfo=localinfo + localinfo=localinfo, + options_prefix=self.options_prefix ) return PetscEq(target, linear_solve) diff --git a/devito/petsc/types/object.py b/devito/petsc/types/object.py index 32892f2aa4..1484e1be81 100644 --- a/devito/petsc/types/object.py +++ b/devito/petsc/types/object.py @@ -1,9 +1,11 @@ from ctypes import POINTER, c_char + from devito.tools import CustomDtype, dtype_to_ctype, as_tuple, CustomIntType from devito.types import (LocalObject, LocalCompositeObject, ModuloDimension, TimeDimension, ArrayObject, CustomDimension) from devito.symbolics import Byref, cast from devito.types.basic import DataSymbol + from devito.petsc.iet.utils import petsc_call diff --git a/devito/petsc/types/types.py b/devito/petsc/types/types.py index 4f235bbbde..baab036f17 100644 --- a/devito/petsc/types/types.py +++ b/devito/petsc/types/types.py @@ -60,7 +60,7 @@ class LinearSolveExpr(MetaData): __rargs__ = ('expr',) __rkwargs__ = ('solver_parameters', 'fielddata', 'time_mapper', - 'localinfo') + 'localinfo', 'options_prefix') defaults = { 'ksp_type': 'gmres', @@ -72,7 +72,8 @@ class LinearSolveExpr(MetaData): } def __new__(cls, expr, solver_parameters=None, - fielddata=None, time_mapper=None, localinfo=None, **kwargs): + fielddata=None, time_mapper=None, localinfo=None, + options_prefix=None, **kwargs): if solver_parameters is None: solver_parameters = cls.defaults @@ -88,6 +89,7 @@ def __new__(cls, expr, solver_parameters=None, obj._fielddata = fielddata if fielddata else FieldData() obj._time_mapper = time_mapper obj._localinfo = localinfo + obj._options_prefix = options_prefix return obj def __repr__(self): @@ -125,6 +127,10 @@ def time_mapper(self): def localinfo(self): return self._localinfo + @property + def options_prefix(self): + return self._options_prefix + @property def grid(self): return self.fielddata.grid diff --git a/devito/symbolics/extended_sympy.py b/devito/symbolics/extended_sympy.py index a75e1091d3..4016a88466 100644 --- a/devito/symbolics/extended_sympy.py +++ b/devito/symbolics/extended_sympy.py @@ -581,7 +581,11 @@ class Keyword(ReservedWord): class String(ReservedWord): - pass + + def __str__(self): + return f'"{self.value}"' + + __repr__ = __str__ class Macro(ReservedWord): diff --git a/tests/test_petsc.py b/tests/test_petsc.py index 3f64a79c82..44668092ad 100644 --- a/tests/test_petsc.py +++ b/tests/test_petsc.py @@ -2,10 +2,12 @@ import numpy as np import os +from collections import OrderedDict from conftest import skipif from devito import (Grid, Function, TimeFunction, Eq, Operator, configuration, norm, switchconfig, SubDomain) +from devito.operator.profiling import PerformanceSummary from devito.ir.iet import (Call, ElementalFunction, FindNodes, retrieve_iteration_tree) from devito.types import Constant, LocalCompositeObject @@ -17,11 +19,11 @@ from devito.petsc.solve import PETScSolve, EssentialBC from devito.petsc.iet.nodes import Expression from devito.petsc.initialize import PetscInitialize +from devito.petsc.logging import PetscSummary -@skipif('petsc') @pytest.fixture(scope='session', autouse=True) -def test_petsc_initialization(): +def petsc_initialization(): # TODO: Temporary workaround until PETSc is automatically # initialized configuration['compiler'] = 'custom' @@ -1377,3 +1379,105 @@ def define(self, dimensions): # Expected norm computed "manually" from sequential run # What rtol and atol should be used? assert np.isclose(norm(u), unorm, rtol=1e-13, atol=1e-13) + + +class TestLogging: + + @skipif('petsc') + @pytest.mark.parametrize('log_level', ['PERF', 'DEBUG']) + def test_logging(self, log_level): + """Verify PetscSummary output when the log level is 'PERF' or 'DEBUG.""" + grid = Grid(shape=(11, 11), dtype=np.float64) + + functions = [Function(name=n, grid=grid, space_order=2) + for n in ['e', 'f']] + e, f = functions + f.data[:] = 5.0 + eq = Eq(e.laplace, f) + + petsc = PETScSolve(eq, target=e, options_prefix='poisson') + + with switchconfig(language='petsc', log_level=log_level): + op = Operator(petsc) + summary = op.apply() + + # One PerformanceSummary + assert len(summary) == 1 + + # Access the PetscSummary + petsc_summary = summary.petsc + + assert isinstance(summary, PerformanceSummary) + assert isinstance(petsc_summary, PetscSummary) + + # One section with a single solver + assert len(petsc_summary) == 1 + + entry0 = petsc_summary.get_entry('section0', 'poisson') + entry1 = petsc_summary[('section0', 'poisson')] + assert entry0 == entry1 + assert str(entry0) == \ + 'PetscEntry(KSPGetIterationNumber=16, SNESGetIterationNumber=1)' + + assert entry0.SNESGetIterationNumber == 1 + + snesiters0 = petsc_summary.SNESGetIterationNumber + snesiters1 = petsc_summary['SNESGetIterationNumber'] + # Check case insensitive key access + snesiters2 = petsc_summary['snesgetiterationnumber'] + snesiters3 = petsc_summary['SNESgetiterationNumber'] + + assert snesiters0 == snesiters1 == snesiters2 == snesiters3 + + assert isinstance(snesiters0, OrderedDict) + assert len(snesiters0) == 1 + key, value = next(iter(snesiters0.items())) + assert str(key) == "PetscKey(name='section0', options_prefix='poisson')" + assert value == 1 + + @skipif('petsc') + def test_logging_multiple_solves(self): + grid = Grid(shape=(11, 11), dtype=np.float64) + + functions = [Function(name=n, grid=grid, space_order=2) + for n in ['e', 'f', 'g', 'h']] + e, f, g, h = functions + + e.data[:] = 5.0 + f.data[:] = 6.0 + + eq1 = Eq(g.laplace, e) + eq2 = Eq(h, f + 5.0) + + solver1 = PETScSolve(eq1, target=g, options_prefix='poisson1') + solver2 = PETScSolve(eq2, target=h, options_prefix='poisson2') + + with switchconfig(language='petsc', log_level='DEBUG'): + op = Operator([solver1, solver2]) + summary = op.apply() + + petsc_summary = summary.petsc + # One PetscKey, PetscEntry for each solver + assert len(petsc_summary) == 2 + + entry1 = petsc_summary.get_entry('section0', 'poisson1') + entry2 = petsc_summary.get_entry('section1', 'poisson2') + + assert str(entry1) == \ + 'PetscEntry(KSPGetIterationNumber=16, SNESGetIterationNumber=1)' + assert str(entry2) == \ + 'PetscEntry(KSPGetIterationNumber=1, SNESGetIterationNumber=1)' + + assert len(petsc_summary.KSPGetIterationNumber) == 2 + assert len(petsc_summary.SNESGetIterationNumber) == 2 + + assert entry1.KSPGetIterationNumber == 16 + assert entry1.SNESGetIterationNumber == 1 + assert entry2.KSPGetIterationNumber == 1 + assert entry2.SNESGetIterationNumber == 1 + + # Test key access to PetscEntry + assert entry1['KSPGetIterationNumber'] == 16 + assert entry1['SNESGetIterationNumber'] == 1 + # Case insensitive key access + assert entry1['kspgetiterationnumber'] == 16 From e3e26f0155be3b7d126ddfd92450775395ba0c46 Mon Sep 17 00:00:00 2001 From: ZoeLeibowitz Date: Mon, 14 Jul 2025 17:27:18 +0100 Subject: [PATCH 73/84] compiler: Fix language summary --- devito/operator/profiling.py | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/devito/operator/profiling.py b/devito/operator/profiling.py index 360d506d17..5e30500761 100644 --- a/devito/operator/profiling.py +++ b/devito/operator/profiling.py @@ -342,11 +342,11 @@ def summary(self, args, dtype, params, reduce_over=None): # data transfers) summary.add_glb_fdlike('fdlike-nosetup', points, reduce_over_nosetup) - # Add the language specific summary - summary.add_language_summary( - self.language, - language_summary_mapper[self.language](params) - ) + # Add the language specific summary if necessary + mapper_func = language_summary_mapper.get(self.language) + if mapper_func: + summary.add_language_summary(self.language, mapper_func(params)) + return summary From 9a182daeb3d914ed92e80fa29b6d265118089a56 Mon Sep 17 00:00:00 2001 From: ZoeLeibowitz Date: Mon, 14 Jul 2025 18:25:22 +0100 Subject: [PATCH 74/84] misc: Fix advisor profiling with new language arg --- devito/operator/profiling.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/devito/operator/profiling.py b/devito/operator/profiling.py index 5e30500761..04f6779904 100644 --- a/devito/operator/profiling.py +++ b/devito/operator/profiling.py @@ -375,11 +375,11 @@ class AdvisorProfiler(AdvancedProfiler): _default_libs = ['ittnotify'] _ext_calls = [_api_resume, _api_pause] - def __init__(self, name): + def __init__(self, name, language): if self._attempted_init: return - super().__init__(name) + super().__init__(name, language) path = get_advisor_path() if path: From 774af10bb6a63861b480369295e56018ce001d10 Mon Sep 17 00:00:00 2001 From: ZoeLeibowitz Date: Wed, 16 Jul 2025 23:05:04 +0100 Subject: [PATCH 75/84] misc: Clean up and docstrings --- devito/petsc/iet/routines.py | 63 ++++++++++++++-------------------- devito/petsc/types/array.py | 2 +- devito/petsc/types/types.py | 35 +++++++++++-------- devito/symbolics/extraction.py | 7 ++-- devito/types/object.py | 2 +- 5 files changed, 54 insertions(+), 55 deletions(-) diff --git a/devito/petsc/iet/routines.py b/devito/petsc/iet/routines.py index d21af4be6b..503d773372 100644 --- a/devito/petsc/iet/routines.py +++ b/devito/petsc/iet/routines.py @@ -42,14 +42,12 @@ def __init__(self, **kwargs): self._struct_params = [] self._main_matvec_callback = None - self._main_formfunc_callback = None self._user_struct_callback = None - # TODO: Test pickling. The mutability of these lists - # could cause issues when pickling? + self._F_efunc = None + self._b_efunc = None + self._J_efuncs = [] - self._F_efuncs = [] - self._b_efuncs = [] - self._initialguesses = [] + self._initial_guesses = [] self._make_core() self._efuncs = self._uxreplace_efuncs() @@ -69,31 +67,26 @@ def filtered_struct_params(self): @property def main_matvec_callback(self): """ - This is the matvec callback associated with the whole Jacobian i.e - is set in the main kernel via - `PetscCall(MatShellSetOperation(J,MATOP_MULT,(void (*)(void))...));` + The matrix-vector callback for the full Jacobian. + This is the function set in the main Kernel via: + PetscCall(MatShellSetOperation(J, MATOP_MULT, (void (*)(void))...)); + The callback has the signature `(Mat, Vec, Vec)`. """ return self._J_efuncs[0] - @property - def main_formfunc_callback(self): - return self._F_efuncs[0] - @property def J_efuncs(self): + """ + List of matrix-vector callbacks. + Each callback has the signature `(Mat, Vec, Vec)`. Typically, this list + contains a single element, but in mixed systems it can include multiple + callbacks, one for each subblock. + """ return self._J_efuncs @property - def F_efuncs(self): - return self._F_efuncs - - @property - def b_efuncs(self): - return self._b_efuncs - - @property - def initialguesses(self): - return self._initialguesses + def initial_guesses(self): + return self._initial_guesses @property def user_struct_callback(self): @@ -284,7 +277,7 @@ def _make_formfunc(self): retval=objs['err'], parameters=(objs['snes'], objs['X'], objs['F'], objs['dummyptr']) ) - self._F_efuncs.append(cb) + self._F_efunc = cb self._efuncs[cb.name] = cb def _create_formfunc_body(self, body): @@ -422,7 +415,7 @@ def _make_formrhs(self): retval=objs['err'], parameters=(sobjs['callbackdm'], objs['B']) ) - self._b_efuncs.append(cb) + self._b_efunc = cb self._efuncs[cb.name] = cb def _create_form_rhs_body(self, body): @@ -534,7 +527,7 @@ def _make_initialguess(self): retval=objs['err'], parameters=(sobjs['callbackdm'], objs['xloc']) ) - self._initialguesses.append(cb) + self._initial_guesses.append(cb) self._efuncs[cb.name] = cb def _create_initial_guess_body(self, body): @@ -660,16 +653,12 @@ def jacobian(self): @property def main_matvec_callback(self): """ - This is the matvec callback associated with the whole Jacobian i.e + This is the matrix-vector callback associated with the whole Jacobian i.e is set in the main kernel via `PetscCall(MatShellSetOperation(J,MATOP_MULT,(void (*)(void))MyMatShellMult));` """ return self._main_matvec_callback - @property - def main_formfunc_callback(self): - return self._main_formfunc_callback - def _make_core(self): for sm in self.fielddata.jacobian.nonzero_submatrices: self._make_matvec(sm, prefix=f'{sm.name}_MatMult') @@ -757,7 +746,7 @@ def _make_whole_formfunc(self): retval=objs['err'], parameters=(objs['snes'], objs['X'], objs['F'], objs['dummyptr']) ) - self._main_formfunc_callback = cb + self._F_efunc = cb self._efuncs[cb.name] = cb def _whole_formfunc_body(self, body): @@ -1310,7 +1299,7 @@ def _setup(self): 'MatShellSetOperation', [sobjs['Jac'], 'MATOP_MULT', MatShellSetOp(matvec.name, void, void)] ) - formfunc = self.cbbuilder.main_formfunc_callback + formfunc = self.cbbuilder._F_efunc formfunc_operation = petsc_call( 'SNESSetFunction', [sobjs['snes'], objs['Null'], FormFunctionCallback(formfunc.name, void, void), @@ -1477,7 +1466,7 @@ def _setup(self): 'MatShellSetOperation', [sobjs['Jac'], 'MATOP_MULT', MatShellSetOp(matvec.name, void, void)] ) - formfunc = self.cbbuilder.main_formfunc_callback + formfunc = self.cbbuilder._F_efunc formfunc_operation = petsc_call( 'SNESSetFunction', [sobjs['snes'], objs['Null'], FormFunctionCallback(formfunc.name, void, void), @@ -1593,7 +1582,7 @@ def _execute_solve(self): struct_assignment = self.timedep.assign_time_iters(sobjs['userctx']) - b_efunc = self.cbbuilder.b_efuncs[0] + b_efunc = self.cbbuilder._b_efunc dmda = sobjs['dmda'] @@ -1601,8 +1590,8 @@ def _execute_solve(self): vec_place_array = self.timedep.place_array(target) - if self.cbbuilder.initialguesses: - initguess = self.cbbuilder.initialguesses[0] + if self.cbbuilder.initial_guesses: + initguess = self.cbbuilder.initial_guesses[0] initguess_call = petsc_call(initguess.name, [dmda, sobjs['xlocal']]) else: initguess_call = None diff --git a/devito/petsc/types/array.py b/devito/petsc/types/array.py index e691a96c33..1bed71ec50 100644 --- a/devito/petsc/types/array.py +++ b/devito/petsc/types/array.py @@ -169,7 +169,7 @@ def _C_ctype(self): fields = [(i.target.name, dtype_to_ctype(i.dtype)) for i in self.components] return POINTER(type(self.pname, (Structure,), {'_fields_': fields})) - @cached_property + @property def symbolic_shape(self): return self.c0.symbolic_shape diff --git a/devito/petsc/types/types.py b/devito/petsc/types/types.py index baab036f17..4f1f492c8e 100644 --- a/devito/petsc/types/types.py +++ b/devito/petsc/types/types.py @@ -1,6 +1,7 @@ import sympy from itertools import chain +from functools import cached_property from devito.tools import Reconstructable, sympy_mutex, as_tuple, frozendict from devito.tools.dtypes_lowering import dtype_mapper @@ -241,7 +242,7 @@ def __init__(self, targets, arrays, jacobian=None, residual=None): self._jacobian = jacobian self._residual = residual - @property + @cached_property def space_dimensions(self): space_dims = {t.space_dimensions for t in self.targets} if len(space_dims) > 1: @@ -251,7 +252,7 @@ def space_dimensions(self): ) return space_dims.pop() - @property + @cached_property def grid(self): """The unique `Grid` associated with all targets.""" grids = [t.grid for t in self.targets] @@ -261,7 +262,7 @@ def grid(self): ) return grids.pop() - @property + @cached_property def space_order(self): # NOTE: since we use DMDA to create vecs for the coupled solves, # all fields must have the same space order @@ -398,7 +399,8 @@ def _build_matvecs(self): matvecs.extend( e for e in self._build_matvec_expr(eq) if e is not None ) - matvecs = tuple(sorted(matvecs, key=lambda e: not isinstance(e, EssentialBC))) + key = lambda e: not isinstance(e, EssentialBC) + matvecs = tuple(sorted(matvecs, key=key)) matvecs = self._scale_non_bcs(matvecs) scdiag = self._compute_scdiag(matvecs) @@ -425,7 +427,7 @@ class MixedJacobian(BaseJacobian): """ def __init__(self, target_exprs, arrays, time_mapper): super().__init__(arrays=arrays, target=None) - self.targets = tuple(target_exprs.keys()) + self.targets = tuple(target_exprs) self.time_mapper = time_mapper self._submatrices = [] self._build_blocks(target_exprs) @@ -443,12 +445,12 @@ def n_submatrices(self): """Return the number of submatrix blocks.""" return len(self._submatrices) - @property + @cached_property def nonzero_submatrices(self): """Return SubMatrixBlock objects that have non-empty matvecs.""" return [m for m in self.submatrices if m.matvecs] - @property + @cached_property def target_scaler_mapper(self): """ Map each row target to the scdiag of its corresponding @@ -467,13 +469,8 @@ def _build_blocks(self, target_exprs): for i, row_target in enumerate(self.targets): exprs = target_exprs[row_target] for j, col_target in enumerate(self.targets): - matvecs = [] - for expr in exprs: - matvecs.extend( - e for e in self._build_matvec_expr( - expr, col_target=col_target, row_target=row_target - ) - ) + + matvecs = self._build_submatrix_matvecs(exprs, row_target, col_target) matvecs = [m for m in matvecs if m is not None] # Sort to put EssentialBC first if any @@ -497,6 +494,16 @@ def _build_blocks(self, target_exprs): ) self._submatrices.append(block) + def _build_submatrix_matvecs(self, exprs, row_target, col_target): + matvecs = [] + for expr in exprs: + matvecs.extend( + e for e in self._build_matvec_expr( + expr, col_target=col_target, row_target=row_target + ) + ) + return matvecs + def get_submatrix(self, row_idx, col_idx): """ Return the SubMatrixBlock at (row_idx, col_idx), or None if not found. diff --git a/devito/symbolics/extraction.py b/devito/symbolics/extraction.py index c0a28433cc..d429f83831 100644 --- a/devito/symbolics/extraction.py +++ b/devito/symbolics/extraction.py @@ -79,10 +79,13 @@ def centre_stencil(expr, target, as_coeff=False): 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. + Parameters ---------- - expr : the expression to extract the centre stencil from - target : the target function whose centre stencil we want + expr : expr-like + The expression to extract the centre stencil from + target : Function + The target function whose centre stencil we want as_coeff : bool, optional If True, return the coefficient of the centre stencil """ diff --git a/devito/types/object.py b/devito/types/object.py index 1153ce81eb..201a087b3d 100644 --- a/devito/types/object.py +++ b/devito/types/object.py @@ -252,10 +252,10 @@ class LocalCompositeObject(CompositeObject, LocalType): __rkwargs__ = ('modifier', 'liveness') def __init__(self, name, pname, fields, modifier=None, liveness='lazy'): - self.modifier = modifier dtype = CustomDtype(f"struct {pname}", modifier=modifier) Object.__init__(self, name, dtype, None) self._pname = pname + self.modifier = modifier assert liveness in ['eager', 'lazy'] self._liveness = liveness self._fields = fields From 910c98731e441a3eebe2e6738354246924516921 Mon Sep 17 00:00:00 2001 From: ZoeLeibowitz Date: Thu, 17 Jul 2025 10:47:56 +0100 Subject: [PATCH 76/84] misc: Clean up --- devito/petsc/clusters.py | 2 +- devito/petsc/iet/logging.py | 4 +- devito/petsc/iet/passes.py | 30 ++++---- devito/petsc/iet/routines.py | 121 ++++++++++++++++----------------- devito/petsc/logging.py | 8 +-- devito/petsc/solve.py | 8 +-- devito/petsc/types/types.py | 24 +++---- devito/symbolics/extraction.py | 4 +- tests/test_petsc.py | 30 ++++---- 9 files changed, 114 insertions(+), 117 deletions(-) diff --git a/devito/petsc/clusters.py b/devito/petsc/clusters.py index e035ccbefc..7171669490 100644 --- a/devito/petsc/clusters.py +++ b/devito/petsc/clusters.py @@ -20,7 +20,7 @@ def petsc_lift(clusters): processed = [] for c in clusters: if isinstance(c.exprs[0].rhs, LinearSolveExpr): - ispace = c.ispace.lift(c.exprs[0].rhs.fielddata.space_dimensions) + ispace = c.ispace.lift(c.exprs[0].rhs.field_data.space_dimensions) processed.append(c.rebuild(ispace=ispace)) else: processed.append(c) diff --git a/devito/petsc/iet/logging.py b/devito/petsc/iet/logging.py index 6a1c74d767..af8b5c4851 100644 --- a/devito/petsc/iet/logging.py +++ b/devito/petsc/iet/logging.py @@ -16,7 +16,7 @@ def __init__(self, level, **kwargs): self.sobjs = kwargs.get('solver_objs') self.sreg = kwargs.get('sregistry') self.section_mapper = kwargs.get('section_mapper', {}) - self.injectsolve = kwargs.get('injectsolve', None) + self.inject_solve = kwargs.get('inject_solve', None) self.function_list = [] @@ -33,7 +33,7 @@ def __init__(self, level, **kwargs): self.statstruct = PetscInfo( name, pname, self.logobjs, self.sobjs, - self.section_mapper, self.injectsolve, + self.section_mapper, self.inject_solve, self.function_list ) diff --git a/devito/petsc/iet/passes.py b/devito/petsc/iet/passes.py index 8d7042bb6c..aadb711c76 100644 --- a/devito/petsc/iet/passes.py +++ b/devito/petsc/iet/passes.py @@ -31,10 +31,10 @@ @iet_pass def lower_petsc(iet, **kwargs): # Check if PETScSolve was used - injectsolve_mapper = MapNodes(Iteration, PetscMetaData, + inject_solve_mapper = MapNodes(Iteration, PetscMetaData, 'groupby').visit(iet) - if not injectsolve_mapper: + if not inject_solve_mapper: return iet, {} if kwargs['language'] not in petsc_languages: @@ -51,7 +51,7 @@ def lower_petsc(iet, **kwargs): if any(filter(lambda i: isinstance(i.expr.rhs, Finalize), data)): return finalize(iet), core_metadata() - unique_grids = {i.expr.rhs.grid for (i,) in injectsolve_mapper.values()} + unique_grids = {i.expr.rhs.grid for (i,) in inject_solve_mapper.values()} # Assumption is that all solves are on the same grid if len(unique_grids) > 1: raise ValueError("All PETScSolves must use the same Grid, but multiple found.") @@ -69,9 +69,9 @@ def lower_petsc(iet, **kwargs): # Map PETScSolve to its Section (for logging) section_mapper = MapNodes(Section, PetscMetaData, 'groupby').visit(iet) - for iters, (injectsolve,) in injectsolve_mapper.items(): + for iters, (inject_solve,) in inject_solve_mapper.items(): - builder = Builder(injectsolve, objs, iters, comm, section_mapper, **kwargs) + builder = Builder(inject_solve, objs, iters, comm, section_mapper, **kwargs) setup.extend(builder.solversetup.calls) @@ -129,18 +129,18 @@ class Builder: and other functionalities as needed. The class will be extended to accommodate different solver types by returning subclasses of the objects initialised in __init__, - depending on the properties of `injectsolve`. + depending on the properties of `inject_solve`. """ - def __init__(self, injectsolve, objs, iters, comm, section_mapper, **kwargs): - self.injectsolve = injectsolve + def __init__(self, inject_solve, objs, iters, comm, section_mapper, **kwargs): + self.inject_solve = inject_solve self.objs = objs self.iters = iters self.comm = comm self.section_mapper = section_mapper self.kwargs = kwargs - self.coupled = isinstance(injectsolve.expr.rhs.fielddata, MultipleFieldData) + self.coupled = isinstance(inject_solve.expr.rhs.field_data, MultipleFieldData) self.common_kwargs = { - 'injectsolve': self.injectsolve, + 'inject_solve': self.inject_solve, 'objs': self.objs, 'iters': self.iters, 'comm': self.comm, @@ -148,7 +148,7 @@ def __init__(self, injectsolve, objs, iters, comm, section_mapper, **kwargs): **self.kwargs } self.common_kwargs['solver_objs'] = self.objbuilder.solver_objs - self.common_kwargs['timedep'] = self.timedep + self.common_kwargs['time_dependence'] = self.time_dependence self.common_kwargs['cbbuilder'] = self.cbbuilder self.common_kwargs['logger'] = self.logger @@ -161,10 +161,10 @@ def objbuilder(self): ) @cached_property - def timedep(self): - time_mapper = self.injectsolve.expr.rhs.time_mapper - timedep_class = TimeDependent if time_mapper else NonTimeDependent - return timedep_class(**self.common_kwargs) + def time_dependence(self): + mapper = self.inject_solve.expr.rhs.time_mapper + time_class = TimeDependent if mapper else NonTimeDependent + return time_class(**self.common_kwargs) @cached_property def cbbuilder(self): diff --git a/devito/petsc/iet/routines.py b/devito/petsc/iet/routines.py index 503d773372..4f7bf01e8e 100644 --- a/devito/petsc/iet/routines.py +++ b/devito/petsc/iet/routines.py @@ -33,10 +33,10 @@ def __init__(self, **kwargs): self.rcompile = kwargs.get('rcompile', None) self.sregistry = kwargs.get('sregistry', None) self.concretize_mapper = kwargs.get('concretize_mapper', {}) - self.timedep = kwargs.get('timedep') + self.time_dependence = kwargs.get('time_dependence') self.objs = kwargs.get('objs') self.solver_objs = kwargs.get('solver_objs') - self.injectsolve = kwargs.get('injectsolve') + self.inject_solve = kwargs.get('inject_solve') self._efuncs = OrderedDict() self._struct_params = [] @@ -93,23 +93,23 @@ def user_struct_callback(self): return self._user_struct_callback @property - def fielddata(self): - return self.injectsolve.expr.rhs.fielddata + def field_data(self): + return self.inject_solve.expr.rhs.field_data @property def arrays(self): - return self.fielddata.arrays + return self.field_data.arrays @property def target(self): - return self.fielddata.target + return self.field_data.target def _make_core(self): - self._make_matvec(self.fielddata.jacobian) + self._make_matvec(self.field_data.jacobian) self._make_formfunc() self._make_formrhs() - if self.fielddata.initialguess.exprs: - self._make_initialguess() + if self.field_data.initial_guess.exprs: + self._make_initial_guess() self._make_user_struct_callback() def _make_matvec(self, jacobian, prefix='MatMult'): @@ -134,7 +134,7 @@ def _make_matvec(self, jacobian, prefix='MatMult'): self._efuncs[cb.name] = cb def _create_matvec_body(self, body, jacobian): - linsolve_expr = self.injectsolve.expr.rhs + linsolve_expr = self.inject_solve.expr.rhs objs = self.objs sobjs = self.solver_objs @@ -145,7 +145,7 @@ def _create_matvec_body(self, body, jacobian): y_matvec = self.arrays[jacobian.row_target]['y'] x_matvec = self.arrays[jacobian.col_target]['x'] - body = self.timedep.uxreplace_time(body) + body = self.time_dependence.uxreplace_time(body) fields = self._dummy_fields(body) @@ -261,7 +261,7 @@ def _create_matvec_body(self, body, jacobian): return body def _make_formfunc(self): - F_exprs = self.fielddata.residual.F_exprs + F_exprs = self.field_data.residual.F_exprs # Compile `F_exprs` into an IET via recursive compilation irs, _ = self.rcompile( F_exprs, options={'mpi': False}, sregistry=self.sregistry, @@ -281,7 +281,7 @@ def _make_formfunc(self): self._efuncs[cb.name] = cb def _create_formfunc_body(self, body): - linsolve_expr = self.injectsolve.expr.rhs + linsolve_expr = self.inject_solve.expr.rhs objs = self.objs sobjs = self.solver_objs arrays = self.arrays @@ -290,7 +290,7 @@ def _create_formfunc_body(self, body): dmda = sobjs['callbackdm'] ctx = objs['dummyctx'] - body = self.timedep.uxreplace_time(body) + body = self.time_dependence.uxreplace_time(body) fields = self._dummy_fields(body) self._struct_params.extend(fields) @@ -397,7 +397,7 @@ def _create_formfunc_body(self, body): return Uxreplace(subs).visit(body) def _make_formrhs(self): - b_exprs = self.fielddata.residual.b_exprs + b_exprs = self.field_data.residual.b_exprs sobjs = self.solver_objs # Compile `b_exprs` into an IET via recursive compilation @@ -419,7 +419,7 @@ def _make_formrhs(self): self._efuncs[cb.name] = cb def _create_form_rhs_body(self, body): - linsolve_expr = self.injectsolve.expr.rhs + linsolve_expr = self.inject_solve.expr.rhs objs = self.objs sobjs = self.solver_objs target = self.target @@ -441,7 +441,7 @@ def _create_form_rhs_body(self, body): sobjs['blocal'] ]) - b_arr = self.fielddata.arrays[target]['b'] + b_arr = self.field_data.arrays[target]['b'] vec_get_array = petsc_call( 'VecGetArray', [sobjs['blocal'], Byref(b_arr._C_symbol)] @@ -451,7 +451,7 @@ def _create_form_rhs_body(self, body): 'DMDAGetLocalInfo', [dmda, Byref(linsolve_expr.localinfo)] ) - body = self.timedep.uxreplace_time(body) + body = self.time_dependence.uxreplace_time(body) fields = self._dummy_fields(body) self._struct_params.extend(fields) @@ -508,8 +508,8 @@ def _create_form_rhs_body(self, body): return Uxreplace(subs).visit(body) - def _make_initialguess(self): - exprs = self.fielddata.initialguess.exprs + def _make_initial_guess(self): + exprs = self.field_data.initial_guess.exprs sobjs = self.solver_objs # Compile initital guess `eqns` into an IET via recursive compilation @@ -531,7 +531,7 @@ def _make_initialguess(self): self._efuncs[cb.name] = cb def _create_initial_guess_body(self, body): - linsolve_expr = self.injectsolve.expr.rhs + linsolve_expr = self.inject_solve.expr.rhs objs = self.objs sobjs = self.solver_objs target = self.target @@ -539,7 +539,7 @@ def _create_initial_guess_body(self, body): dmda = sobjs['callbackdm'] ctx = objs['dummyctx'] - x_arr = self.fielddata.arrays[target]['x'] + x_arr = self.field_data.arrays[target]['x'] vec_get_array = petsc_call( 'VecGetArray', [objs['xloc'], Byref(x_arr._C_symbol)] @@ -549,7 +549,7 @@ def _create_initial_guess_body(self, body): 'DMDAGetLocalInfo', [dmda, Byref(linsolve_expr.localinfo)] ) - body = self.timedep.uxreplace_time(body) + body = self.time_dependence.uxreplace_time(body) fields = self._dummy_fields(body) self._struct_params.extend(fields) @@ -648,7 +648,7 @@ def submatrices_callback(self): @property def jacobian(self): - return self.injectsolve.expr.rhs.fielddata.jacobian + return self.inject_solve.expr.rhs.field_data.jacobian @property def main_matvec_callback(self): @@ -660,7 +660,7 @@ def main_matvec_callback(self): return self._main_matvec_callback def _make_core(self): - for sm in self.fielddata.jacobian.nonzero_submatrices: + for sm in self.field_data.jacobian.nonzero_submatrices: self._make_matvec(sm, prefix=f'{sm.name}_MatMult') self._make_whole_matvec() @@ -731,7 +731,7 @@ def _whole_matvec_body(self): ) def _make_whole_formfunc(self): - F_exprs = self.fielddata.residual.F_exprs + F_exprs = self.field_data.residual.F_exprs # Compile `F_exprs` into an IET via recursive compilation irs, _ = self.rcompile( F_exprs, options={'mpi': False}, sregistry=self.sregistry, @@ -750,14 +750,14 @@ def _make_whole_formfunc(self): self._efuncs[cb.name] = cb def _whole_formfunc_body(self, body): - linsolve_expr = self.injectsolve.expr.rhs + linsolve_expr = self.inject_solve.expr.rhs objs = self.objs sobjs = self.solver_objs dmda = sobjs['callbackdm'] ctx = objs['dummyctx'] - body = self.timedep.uxreplace_time(body) + body = self.time_dependence.uxreplace_time(body) fields = self._dummy_fields(body) self._struct_params.extend(fields) @@ -1026,11 +1026,11 @@ class BaseObjectBuilder: method to support specific use cases. """ def __init__(self, **kwargs): - self.injectsolve = kwargs.get('injectsolve') + self.inject_solve = kwargs.get('inject_solve') self.objs = kwargs.get('objs') self.sregistry = kwargs.get('sregistry') self.comm = kwargs.get('comm') - self.fielddata = self.injectsolve.expr.rhs.fielddata + self.field_data = self.inject_solve.expr.rhs.field_data self.solver_objs = self._build() def _build(self): @@ -1055,10 +1055,10 @@ def _build(self): functions via `SNESGetDM`. """ sreg = self.sregistry - targets = self.fielddata.targets + targets = self.field_data.targets snes_name = sreg.make_name(prefix='snes') - options_prefix = self.injectsolve.expr.rhs.options_prefix + options_prefix = self.inject_solve.expr.rhs.options_prefix base_dict = { 'Jac': Mat(sreg.make_name(prefix='J')), @@ -1085,7 +1085,7 @@ def _target_dependent(self, base_dict): that will be updated at each time step. """ sreg = self.sregistry - target = self.fielddata.target + target = self.field_data.target base_dict[f'{target.name}_ptr'] = StartPtr( sreg.make_name(prefix=f'{target.name}_ptr'), target.dtype ) @@ -1102,8 +1102,8 @@ class CoupledObjectBuilder(BaseObjectBuilder): def _extend_build(self, base_dict): sreg = self.sregistry objs = self.objs - targets = self.fielddata.targets - arrays = self.fielddata.arrays + targets = self.field_data.targets + arrays = self.field_data.arrays base_dict['fields'] = PointerIS( name=sreg.make_name(prefix='fields'), nindices=len(targets) @@ -1113,14 +1113,14 @@ def _extend_build(self, base_dict): ) base_dict['nfields'] = PetscInt(sreg.make_name(prefix='nfields')) - space_dims = len(self.fielddata.grid.dimensions) + space_dims = len(self.field_data.grid.dimensions) dim_labels = ["M", "N", "P"] base_dict.update({ dim_labels[i]: PetscInt(dim_labels[i]) for i in range(space_dims) }) - submatrices = self.fielddata.jacobian.nonzero_submatrices + submatrices = self.field_data.jacobian.nonzero_submatrices base_dict['jacctx'] = JacobianStruct( name=sreg.make_name(prefix=objs['ljacctx'].name), @@ -1174,7 +1174,7 @@ def _extend_build(self, base_dict): def _target_dependent(self, base_dict): sreg = self.sregistry - targets = self.fielddata.targets + targets = self.field_data.targets for t in targets: name = t.name base_dict[f'{name}_ptr'] = StartPtr( @@ -1208,11 +1208,11 @@ def _target_dependent(self, base_dict): class BaseSetup: def __init__(self, **kwargs): - self.injectsolve = kwargs.get('injectsolve') + self.inject_solve = kwargs.get('inject_solve') self.objs = kwargs.get('objs') self.solver_objs = kwargs.get('solver_objs') self.cbbuilder = kwargs.get('cbbuilder') - self.fielddata = self.injectsolve.expr.rhs.fielddata + self.field_data = self.inject_solve.expr.rhs.field_data self.calls = self._setup() @property @@ -1229,7 +1229,7 @@ def _setup(self): dmda = sobjs['dmda'] - solver_params = self.injectsolve.expr.rhs.solver_parameters + solver_params = self.inject_solve.expr.rhs.solver_parameters snes_create = petsc_call('SNESCreate', [sobjs['comm'], Byref(sobjs['snes'])]) @@ -1252,7 +1252,7 @@ def _setup(self): global_x = petsc_call('DMCreateGlobalVector', [dmda, Byref(sobjs['xglobal'])]) - target = self.fielddata.target + target = self.field_data.target field_from_ptr = FieldFromPointer( target.function._C_field_data, target.function._C_symbol ) @@ -1366,7 +1366,7 @@ def _create_dmda_calls(self, dmda): def _create_dmda(self, dmda): objs = self.objs sobjs = self.solver_objs - grid = self.fielddata.grid + grid = self.field_data.grid nspace_dims = len(grid.dimensions) # MPI communicator @@ -1389,8 +1389,8 @@ def _create_dmda(self, dmda): args.append(dmda.dofs) # "Stencil width" -> size of overlap # TODO: Instead, this probably should be - # extracted from fielddata.target._size_outhalo? - stencil_width = self.fielddata.space_order + # extracted from field_data.target._size_outhalo? + stencil_width = self.field_data.space_order args.append(stencil_width) args.extend([objs['Null']]*nspace_dims) @@ -1411,7 +1411,7 @@ def _setup(self): sobjs = self.solver_objs dmda = sobjs['dmda'] - solver_params = self.injectsolve.expr.rhs.solver_parameters + solver_params = self.inject_solve.expr.rhs.solver_parameters snes_create = petsc_call('SNESCreate', [sobjs['comm'], Byref(sobjs['snes'])]) @@ -1518,7 +1518,7 @@ def _setup(self): Byref(FieldFromComposite(objs['Submats'].base, sobjs['jacctx']))] ) - targets = self.fielddata.targets + targets = self.field_data.targets deref_dms = [ DummyExpr(sobjs[f'da{t.name}'], sobjs['subdms'].indexed[i]) @@ -1563,13 +1563,12 @@ def _setup(self): class Solver: def __init__(self, **kwargs): - self.injectsolve = kwargs.get('injectsolve') + self.inject_solve = kwargs.get('inject_solve') self.objs = kwargs.get('objs') self.solver_objs = kwargs.get('solver_objs') self.iters = kwargs.get('iters') self.cbbuilder = kwargs.get('cbbuilder') - self.timedep = kwargs.get('timedep') - # TODO: Should/could _execute_solve be a cached_property? + self.time_dependence = kwargs.get('time_dependence') self.calls = self._execute_solve() def _execute_solve(self): @@ -1578,9 +1577,9 @@ def _execute_solve(self): the necessary calls to execute the SNES solver. """ sobjs = self.solver_objs - target = self.injectsolve.expr.rhs.fielddata.target + target = self.inject_solve.expr.rhs.field_data.target - struct_assignment = self.timedep.assign_time_iters(sobjs['userctx']) + struct_assignment = self.time_dependence.assign_time_iters(sobjs['userctx']) b_efunc = self.cbbuilder._b_efunc @@ -1588,7 +1587,7 @@ def _execute_solve(self): rhs_call = petsc_call(b_efunc.name, [sobjs['dmda'], sobjs['bglobal']]) - vec_place_array = self.timedep.place_array(target) + vec_place_array = self.time_dependence.place_array(target) if self.cbbuilder.initial_guesses: initguess = self.cbbuilder.initial_guesses[0] @@ -1609,7 +1608,7 @@ def _execute_solve(self): dmda, sobjs['xglobal'], insert_vals, sobjs['xlocal']] ) - vec_reset_array = self.timedep.reset_array(target) + vec_reset_array = self.time_dependence.reset_array(target) run_solver_calls = (struct_assignment,) + ( rhs_call, @@ -1631,7 +1630,7 @@ def spatial_body(self): root = filter_iterations(tree, key=lambda i: i.dim.is_Space) if root: root = root[0] - if self.injectsolve in FindNodes(PetscMetaData).visit(root): + if self.inject_solve in FindNodes(PetscMetaData).visit(root): spatial_body.append(root) spatial_body, = spatial_body return spatial_body @@ -1647,8 +1646,8 @@ def _execute_solve(self): sobjs = self.solver_objs xglob = sobjs['xglobal'] - struct_assignment = self.timedep.assign_time_iters(sobjs['userctx']) - targets = self.injectsolve.expr.rhs.fielddata.targets + struct_assignment = self.time_dependence.assign_time_iters(sobjs['userctx']) + targets = self.inject_solve.expr.rhs.field_data.targets # TODO: optimise the ccode generated here pre_solve = () @@ -1667,7 +1666,7 @@ def _execute_solve(self): petsc_call('DMCreateLocalVector', [dm, Byref(target_xloc)]), # TODO: Need to call reset array - self.timedep.place_array(t), + self.time_dependence.place_array(t), petsc_call( 'DMLocalToGlobal', [dm, target_xloc, insert_vals, target_xglob] @@ -1709,12 +1708,12 @@ def _execute_solve(self): class NonTimeDependent: def __init__(self, **kwargs): - self.injectsolve = kwargs.get('injectsolve') + self.inject_solve = kwargs.get('inject_solve') self.iters = kwargs.get('iters') self.sobjs = kwargs.get('solver_objs') self.kwargs = kwargs self.origin_to_moddim = self._origin_to_moddim_mapper(self.iters) - self.time_idx_to_symb = self.injectsolve.expr.rhs.time_mapper + self.time_idx_to_symb = self.inject_solve.expr.rhs.time_mapper def _origin_to_moddim_mapper(self, iters): return {} @@ -1773,7 +1772,7 @@ class TimeDependent(NonTimeDependent): """ @property def time_spacing(self): - return self.injectsolve.expr.rhs.grid.stepping_dim.spacing + return self.inject_solve.expr.rhs.grid.stepping_dim.spacing @cached_property def symb_to_moddim(self): diff --git a/devito/petsc/logging.py b/devito/petsc/logging.py index f552b69969..40ad2711fa 100644 --- a/devito/petsc/logging.py +++ b/devito/petsc/logging.py @@ -120,15 +120,15 @@ def __getitem__(self, key): class PetscInfo(CompositeObject): __rargs__ = ('name', 'pname', 'logobjs', 'sobjs', 'section_mapper', - 'injectsolve', 'function_list') + 'inject_solve', 'function_list') def __init__(self, name, pname, logobjs, sobjs, section_mapper, - injectsolve, function_list): + inject_solve, function_list): self.logobjs = logobjs self.sobjs = sobjs self.section_mapper = section_mapper - self.injectsolve = injectsolve + self.inject_solve = inject_solve self.function_list = function_list mapper = {v: k for k, v in petsc_type_mappings.items()} @@ -138,7 +138,7 @@ def __init__(self, name, pname, logobjs, sobjs, section_mapper, @property def section(self): section = self.section_mapper.items() - return next((k[0].name for k, v in section if self.injectsolve in v), None) + return next((k[0].name for k, v in section if self.inject_solve in v), None) @property def summary_key(self): diff --git a/devito/petsc/solve.py b/devito/petsc/solve.py index 55633d40fd..77022fc6b8 100644 --- a/devito/petsc/solve.py +++ b/devito/petsc/solve.py @@ -68,13 +68,13 @@ def __init__(self, solver_parameters=None, target_exprs=None, options_prefix=Non self.options_prefix = options_prefix def build_expr(self): - target, funcs, fielddata = self.linear_solve_args() + target, funcs, field_data = self.linear_solve_args() # Placeholder expression for inserting calls to the solver linear_solve = LinearSolveExpr( funcs, self.solver_params, - fielddata=fielddata, + field_data=field_data, time_mapper=self.time_mapper, localinfo=localinfo, options_prefix=self.options_prefix @@ -93,13 +93,13 @@ def linear_solve_args(self): jacobian = Jacobian(target, exprs, arrays, self.time_mapper) residual = Residual(target, exprs, arrays, self.time_mapper, jacobian.scdiag) - initialguess = InitialGuess(target, exprs, arrays) + initial_guess = InitialGuess(target, exprs, arrays) field_data = FieldData( target=target, jacobian=jacobian, residual=residual, - initialguess=initialguess, + initial_guess=initial_guess, arrays=arrays ) diff --git a/devito/petsc/types/types.py b/devito/petsc/types/types.py index 4f1f492c8e..fa87920571 100644 --- a/devito/petsc/types/types.py +++ b/devito/petsc/types/types.py @@ -60,7 +60,7 @@ class LinearSolveExpr(MetaData): """ __rargs__ = ('expr',) - __rkwargs__ = ('solver_parameters', 'fielddata', 'time_mapper', + __rkwargs__ = ('solver_parameters', 'field_data', 'time_mapper', 'localinfo', 'options_prefix') defaults = { @@ -73,7 +73,7 @@ class LinearSolveExpr(MetaData): } def __new__(cls, expr, solver_parameters=None, - fielddata=None, time_mapper=None, localinfo=None, + field_data=None, time_mapper=None, localinfo=None, options_prefix=None, **kwargs): if solver_parameters is None: @@ -87,7 +87,7 @@ def __new__(cls, expr, solver_parameters=None, obj._expr = expr obj._solver_parameters = solver_parameters - obj._fielddata = fielddata if fielddata else FieldData() + obj._field_data = field_data if field_data else FieldData() obj._time_mapper = time_mapper obj._localinfo = localinfo obj._options_prefix = options_prefix @@ -113,8 +113,8 @@ def expr(self): return self._expr @property - def fielddata(self): - return self._fielddata + def field_data(self): + return self._field_data @property def solver_parameters(self): @@ -134,7 +134,7 @@ def options_prefix(self): @property def grid(self): - return self.fielddata.grid + return self.field_data.grid @classmethod def eval(cls, *args): @@ -158,14 +158,14 @@ class FieldData: the PETScArray representing the `target`. residual : Residual Defines the residual function F(target) = 0. - initialguess : InitialGuess + initial_guess : InitialGuess Defines the initial guess for the solution, which satisfies essential boundary conditions. arrays : dict A dictionary mapping `target` to its corresponding PETScArrays. """ def __init__(self, target=None, jacobian=None, residual=None, - initialguess=None, arrays=None, **kwargs): + initial_guess=None, arrays=None, **kwargs): self._target = target petsc_precision = dtype_mapper[petsc_variables['PETSC_PRECISION']] if self._target.dtype != petsc_precision: @@ -176,7 +176,7 @@ def __init__(self, target=None, jacobian=None, residual=None, ) self._jacobian = jacobian self._residual = residual - self._initialguess = initialguess + self._initial_guess = initial_guess self._arrays = arrays @property @@ -192,8 +192,8 @@ def residual(self): return self._residual @property - def initialguess(self): - return self._initialguess + def initial_guess(self): + return self._initial_guess @property def arrays(self): @@ -230,7 +230,7 @@ class MultipleFieldData(FieldData): Defines the matrix-vector products for the full system Jacobian. residual : MixedResidual Defines the residual function F(targets) = 0. - initialguess : InitialGuess + initial_guess : InitialGuess Defines the initial guess metadata, which satisfies essential boundary conditions. arrays : dict diff --git a/devito/symbolics/extraction.py b/devito/symbolics/extraction.py index d429f83831..45a48c80e8 100644 --- a/devito/symbolics/extraction.py +++ b/devito/symbolics/extraction.py @@ -76,9 +76,7 @@ def _(expr, targets): @singledispatch def centre_stencil(expr, target, as_coeff=False): """ - 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. + Extract the centre stencil from an expression. Parameters ---------- diff --git a/tests/test_petsc.py b/tests/test_petsc.py index 44668092ad..f297b7993f 100644 --- a/tests/test_petsc.py +++ b/tests/test_petsc.py @@ -214,7 +214,7 @@ def test_LinearSolveExpr(): eqn = Eq(f, g.laplace) - linsolveexpr = LinearSolveExpr(eqn.rhs, fielddata=FieldData(target=f)) + linsolveexpr = LinearSolveExpr(eqn.rhs, field_data=FieldData(target=f)) # Check the solver parameters assert linsolveexpr.solver_parameters == \ @@ -653,7 +653,7 @@ def define(self, dimensions): petsc = PETScSolve([eq1, bc_1, bc_2], target=e) - jac = petsc.rhs.fielddata.jacobian + jac = petsc.rhs.field_data.jacobian assert jac.row_target == e assert jac.col_target == e @@ -699,7 +699,7 @@ def define(self, dimensions): petsc = PETScSolve([eq1, bc_1, bc_2], target=e) - res = petsc.rhs.fielddata.residual + res = petsc.rhs.field_data.residual assert res.target == e # NOTE: This is likely to change when PetscSection + DMDA is supported @@ -789,14 +789,14 @@ def test_coupled_vs_non_coupled(self, eq1, eq2, so): assert len(callbacks1) == 8 assert len(callbacks2) == 6 - # Check fielddata type - fielddata1 = petsc1.rhs.fielddata - fielddata2 = petsc2.rhs.fielddata - fielddata3 = petsc3.rhs.fielddata + # Check field_data type + field0 = petsc1.rhs.field_data + field1 = petsc2.rhs.field_data + field2 = petsc3.rhs.field_data - assert isinstance(fielddata1, FieldData) - assert isinstance(fielddata2, FieldData) - assert isinstance(fielddata3, MultipleFieldData) + assert isinstance(field0, FieldData) + assert isinstance(field1, FieldData) + assert isinstance(field2, MultipleFieldData) @skipif('petsc') def test_coupled_structs(self): @@ -919,7 +919,7 @@ def test_mixed_jacobian(self): petsc = PETScSolve({e: [eq1], g: [eq2]}) - jacobian = petsc.rhs.fielddata.jacobian + jacobian = petsc.rhs.field_data.jacobian j00 = jacobian.get_submatrix(0, 0) j01 = jacobian.get_submatrix(0, 1) @@ -1031,7 +1031,7 @@ def test_coupling(self, eq1, eq2, j01_matvec, j10_matvec): petsc = PETScSolve({e: [eq1], g: [eq2]}) - jacobian = petsc.rhs.fielddata.jacobian + jacobian = petsc.rhs.field_data.jacobian j01 = jacobian.get_submatrix(0, 1) j10 = jacobian.get_submatrix(1, 0) @@ -1080,7 +1080,7 @@ def test_jacobian_scaling_1D(self, eq1, eq2, so, scale): petsc = PETScSolve({e: [eq1], g: [eq2]}) - jacobian = petsc.rhs.fielddata.jacobian + jacobian = petsc.rhs.field_data.jacobian j00 = jacobian.get_submatrix(0, 0) j11 = jacobian.get_submatrix(1, 1) @@ -1130,7 +1130,7 @@ def test_jacobian_scaling_2D(self, eq1, eq2, so, scale): petsc = PETScSolve({e: [eq1], g: [eq2]}) - jacobian = petsc.rhs.fielddata.jacobian + jacobian = petsc.rhs.field_data.jacobian j00 = jacobian.get_submatrix(0, 0) j11 = jacobian.get_submatrix(1, 1) @@ -1183,7 +1183,7 @@ def test_jacobian_scaling_3D(self, eq1, eq2, so, scale): petsc = PETScSolve({e: [eq1], g: [eq2]}) - jacobian = petsc.rhs.fielddata.jacobian + jacobian = petsc.rhs.field_data.jacobian j00 = jacobian.get_submatrix(0, 0) j11 = jacobian.get_submatrix(1, 1) From ff4b2d7185e939b8bdc7a3e45be5ad37c008d181 Mon Sep 17 00:00:00 2001 From: ZoeLeibowitz Date: Thu, 17 Jul 2025 11:28:08 +0100 Subject: [PATCH 77/84] misc: flake8 --- devito/petsc/iet/passes.py | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/devito/petsc/iet/passes.py b/devito/petsc/iet/passes.py index aadb711c76..d0118d6863 100644 --- a/devito/petsc/iet/passes.py +++ b/devito/petsc/iet/passes.py @@ -32,7 +32,7 @@ def lower_petsc(iet, **kwargs): # Check if PETScSolve was used inject_solve_mapper = MapNodes(Iteration, PetscMetaData, - 'groupby').visit(iet) + 'groupby').visit(iet) if not inject_solve_mapper: return iet, {} @@ -73,7 +73,7 @@ def lower_petsc(iet, **kwargs): builder = Builder(inject_solve, objs, iters, comm, section_mapper, **kwargs) - setup.extend(builder.solversetup.calls) + setup.extend(builder.solver_setup.calls) # Transform the spatial iteration loop with the calls to execute the solver subs.update({builder.solve.spatial_body: builder.calls}) @@ -147,13 +147,13 @@ def __init__(self, inject_solve, objs, iters, comm, section_mapper, **kwargs): 'section_mapper': self.section_mapper, **self.kwargs } - self.common_kwargs['solver_objs'] = self.objbuilder.solver_objs + self.common_kwargs['solver_objs'] = self.object_builder.solver_objs self.common_kwargs['time_dependence'] = self.time_dependence self.common_kwargs['cbbuilder'] = self.cbbuilder self.common_kwargs['logger'] = self.logger @cached_property - def objbuilder(self): + def object_builder(self): return ( CoupledObjectBuilder(**self.common_kwargs) if self.coupled else @@ -172,7 +172,7 @@ def cbbuilder(self): if self.coupled else CBBuilder(**self.common_kwargs) @cached_property - def solversetup(self): + def solver_setup(self): return CoupledSetup(**self.common_kwargs) \ if self.coupled else BaseSetup(**self.common_kwargs) From 3044691c0ac7397a80842c2f29f35c0e419b7d64 Mon Sep 17 00:00:00 2001 From: ZoeLeibowitz Date: Mon, 18 Aug 2025 17:12:56 +0100 Subject: [PATCH 78/84] tests: Update petsc tests since merging with latest main --- tests/test_petsc.py | 28 ++++++++++++++-------------- 1 file changed, 14 insertions(+), 14 deletions(-) diff --git a/tests/test_petsc.py b/tests/test_petsc.py index f297b7993f..207294dbaa 100644 --- a/tests/test_petsc.py +++ b/tests/test_petsc.py @@ -247,15 +247,15 @@ def test_dmda_create(): op3 = Operator(petsc3, opt='noop') assert 'PetscCall(DMDACreate1d(PETSC_COMM_WORLD,DM_BOUNDARY_GHOSTED,' + \ - '2,1,2,NULL,&(da0)));' in str(op1) + '2,1,2,NULL,&da0));' in str(op1) assert 'PetscCall(DMDACreate2d(PETSC_COMM_WORLD,DM_BOUNDARY_GHOSTED,' + \ - 'DM_BOUNDARY_GHOSTED,DMDA_STENCIL_BOX,2,2,1,1,1,4,NULL,NULL,&(da0)));' \ + 'DM_BOUNDARY_GHOSTED,DMDA_STENCIL_BOX,2,2,1,1,1,4,NULL,NULL,&da0));' \ in str(op2) assert 'PetscCall(DMDACreate3d(PETSC_COMM_WORLD,DM_BOUNDARY_GHOSTED,' + \ 'DM_BOUNDARY_GHOSTED,DM_BOUNDARY_GHOSTED,DMDA_STENCIL_BOX,6,5,4' + \ - ',1,1,1,1,6,NULL,NULL,NULL,&(da0)));' in str(op3) + ',1,1,1,1,6,NULL,NULL,NULL,&da0));' in str(op3) @skipif('petsc') @@ -390,12 +390,12 @@ def test_petsc_frees(): frees = op.body.frees # Check the frees appear in the following order - assert str(frees[0]) == 'PetscCall(VecDestroy(&(bglobal0)));' - assert str(frees[1]) == 'PetscCall(VecDestroy(&(xglobal0)));' - assert str(frees[2]) == 'PetscCall(VecDestroy(&(xlocal0)));' - assert str(frees[3]) == 'PetscCall(MatDestroy(&(J0)));' - assert str(frees[4]) == 'PetscCall(SNESDestroy(&(snes0)));' - assert str(frees[5]) == 'PetscCall(DMDestroy(&(da0)));' + assert str(frees[0]) == 'PetscCall(VecDestroy(&bglobal0));' + assert str(frees[1]) == 'PetscCall(VecDestroy(&xglobal0));' + assert str(frees[2]) == 'PetscCall(VecDestroy(&xlocal0));' + assert str(frees[3]) == 'PetscCall(MatDestroy(&J0));' + assert str(frees[4]) == 'PetscCall(SNESDestroy(&snes0));' + assert str(frees[5]) == 'PetscCall(DMDestroy(&da0));' @skipif('petsc') @@ -863,13 +863,13 @@ def test_coupled_frees(self, n_fields): # IS Destroy calls for i in range(n_fields): - assert str(frees[i]) == f'PetscCall(ISDestroy(&(fields0[{i}])));' + assert str(frees[i]) == f'PetscCall(ISDestroy(&fields0[{i}]));' assert str(frees[n_fields]) == 'PetscCall(PetscFree(fields0));' # DM Destroy calls for i in range(n_fields): assert str(frees[n_fields + 1 + i]) == \ - f'PetscCall(DMDestroy(&(subdms0[{i}])));' + f'PetscCall(DMDestroy(&subdms0[{i}]));' assert str(frees[n_fields*2 + 1]) == 'PetscCall(PetscFree(subdms0));' @skipif('petsc') @@ -895,15 +895,15 @@ def test_dmda_dofs(self): # Check the number of dofs in the DMDA for each field assert 'PetscCall(DMDACreate2d(PETSC_COMM_WORLD,DM_BOUNDARY_GHOSTED,' + \ - 'DM_BOUNDARY_GHOSTED,DMDA_STENCIL_BOX,11,11,1,1,1,2,NULL,NULL,&(da0)));' \ + 'DM_BOUNDARY_GHOSTED,DMDA_STENCIL_BOX,11,11,1,1,1,2,NULL,NULL,&da0));' \ in str(op1) assert 'PetscCall(DMDACreate2d(PETSC_COMM_WORLD,DM_BOUNDARY_GHOSTED,' + \ - 'DM_BOUNDARY_GHOSTED,DMDA_STENCIL_BOX,11,11,1,1,2,2,NULL,NULL,&(da0)));' \ + 'DM_BOUNDARY_GHOSTED,DMDA_STENCIL_BOX,11,11,1,1,2,2,NULL,NULL,&da0));' \ in str(op2) assert 'PetscCall(DMDACreate2d(PETSC_COMM_WORLD,DM_BOUNDARY_GHOSTED,' + \ - 'DM_BOUNDARY_GHOSTED,DMDA_STENCIL_BOX,11,11,1,1,3,2,NULL,NULL,&(da0)));' \ + 'DM_BOUNDARY_GHOSTED,DMDA_STENCIL_BOX,11,11,1,1,3,2,NULL,NULL,&da0));' \ in str(op3) @skipif('petsc') From 4139a3f75765d9efc8d1f8f19f1146dc5c7f6b8c Mon Sep 17 00:00:00 2001 From: ZoeLeibowitz Date: Mon, 18 Aug 2025 17:19:49 +0100 Subject: [PATCH 79/84] misc: Flake8 --- tests/test_symbolics.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/test_symbolics.py b/tests/test_symbolics.py index b41c2629c9..cda19cb3ed 100644 --- a/tests/test_symbolics.py +++ b/tests/test_symbolics.py @@ -16,7 +16,7 @@ DefFunction, FieldFromPointer, INT, FieldFromComposite, IntDiv, Namespace, Rvalue, ReservedWord, ListInitializer, uxreplace, pow_to_mul, retrieve_derivatives, BaseCast, SizeOf, VectorAccess, separate_eqn, - centre_stencil + centre_stencil, sympy_dtype ) from devito.tools import as_tuple, CustomDtype from devito.types import (Array, Bundle, FIndexed, LocalObject, Object, From 7c69c4dcb5c52c66751adbabbfd811b0940565ab Mon Sep 17 00:00:00 2001 From: ZoeLeibowitz <90093761+ZoeLeibowitz@users.noreply.github.com> Date: Fri, 19 Sep 2025 10:00:14 +0100 Subject: [PATCH 80/84] dsl/compiler: Extend solver_parameters handling in PETScSolve and update logging (#2718) * compiler: Utilise petsctools and start extending solver parameters functionality examples: Add example using argparse in conjunction with PetscInitialize * compiler: Edit argument to PetscInitialize * compiler: Use petsctools to process solver params and start callback to set petsc options * compiler: Progress on petscoptions callbacks * dsl/compiler: Re-factor solver params, add solver_parameters.py file * misc: Add tests * compiler: Start extending the PetscSummary * dsl: Fix hashing for solveexpr * misc: Add tests and clean up * compiler: Add utility function inside petsc routines * misc: Clean up * tests: Add command line tests with random prefixes * compiler: Create dummy op for petscgetargs * misc: Update requirements * misc: Create getargs Op inside function * compiler/dsl: Add get_info functionality to petscsolve * compiler/misc: Drop petscgetargs callback, add functions to logging/get_info, clean up and more tests * workflows: Fix serial wf run --- .github/workflows/pytest-petsc.yml | 8 +- devito/operator/profiling.py | 7 +- devito/petsc/clusters.py | 4 +- devito/petsc/iet/logging.py | 77 ++-- devito/petsc/iet/passes.py | 49 +- devito/petsc/iet/routines.py | 312 ++++++------- devito/petsc/iet/utils.py | 3 +- devito/petsc/initialize.py | 26 +- devito/petsc/logging.py | 142 ++++-- devito/petsc/solve.py | 112 +++-- devito/petsc/solver_parameters.py | 41 ++ devito/petsc/types/object.py | 16 + devito/petsc/types/types.py | 146 +++--- devito/petsc/utils.py | 78 +++- devito/symbolics/extended_sympy.py | 3 + requirements.txt | 3 +- tests/test_petsc.py | 689 +++++++++++++++++++++++++++-- 17 files changed, 1244 insertions(+), 472 deletions(-) create mode 100644 devito/petsc/solver_parameters.py diff --git a/.github/workflows/pytest-petsc.yml b/.github/workflows/pytest-petsc.yml index 502e9b2139..b3ca034500 100644 --- a/.github/workflows/pytest-petsc.yml +++ b/.github/workflows/pytest-petsc.yml @@ -72,9 +72,13 @@ jobs: 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 + - name: Test with pytest - serial run: | - ${{ env.RUN_CMD }} mpiexec -n 1 pytest --cov --cov-config=.coveragerc --cov-report=xml ${{ env.TESTS }} + ${{ env.RUN_CMD }} mpiexec -n 1 pytest -m "not parallel" --cov --cov-config=.coveragerc --cov-report=xml ${{ env.TESTS }} + + - name: Test with pytest - parallel + run: | + ${{ env.RUN_CMD }} python3 -m pytest --cov --cov-config=.coveragerc --cov-report=xml -m parallel ${{ env.TESTS }} - name: Test examples run: | diff --git a/devito/operator/profiling.py b/devito/operator/profiling.py index 349a09f0b4..66c45052d6 100644 --- a/devito/operator/profiling.py +++ b/devito/operator/profiling.py @@ -213,6 +213,11 @@ def summary(self, args, dtype, params, reduce_over=None): else: summary.add(name, None, time) + # Add the language-specific summary if necessary + mapper_func = language_summary_mapper.get(self.language) + if mapper_func: + summary.add_language_summary(self.language, mapper_func(params)) + return summary @@ -342,7 +347,7 @@ def summary(self, args, dtype, params, reduce_over=None): # data transfers) summary.add_glb_fdlike('fdlike-nosetup', points, reduce_over_nosetup) - # Add the language specific summary if necessary + # Add the language-specific summary if necessary mapper_func = language_summary_mapper.get(self.language) if mapper_func: summary.add_language_summary(self.language, mapper_func(params)) diff --git a/devito/petsc/clusters.py b/devito/petsc/clusters.py index 7171669490..00d3e450b7 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 SolverMetaData @timed_pass() @@ -19,7 +19,7 @@ def petsc_lift(clusters): """ processed = [] for c in clusters: - if isinstance(c.exprs[0].rhs, LinearSolveExpr): + if isinstance(c.exprs[0].rhs, SolverMetaData): ispace = c.ispace.lift(c.exprs[0].rhs.field_data.space_dimensions) processed.append(c.rebuild(ispace=ispace)) else: diff --git a/devito/petsc/iet/logging.py b/devito/petsc/iet/logging.py index af8b5c4851..68f31cfa91 100644 --- a/devito/petsc/iet/logging.py +++ b/devito/petsc/iet/logging.py @@ -3,6 +3,7 @@ from devito.symbolics import Byref, FieldFromPointer from devito.ir.iet import DummyExpr from devito.logger import PERF +from devito.tools import frozendict from devito.petsc.iet.utils import petsc_call from devito.petsc.logging import petsc_return_variable_dict, PetscInfo @@ -12,19 +13,29 @@ class PetscLogger: """ Class for PETSc loggers that collect solver related statistics. """ + # TODO: Update docstring with kwargs def __init__(self, level, **kwargs): + + self.query_functions = kwargs.get('get_info', []) self.sobjs = kwargs.get('solver_objs') self.sreg = kwargs.get('sregistry') self.section_mapper = kwargs.get('section_mapper', {}) self.inject_solve = kwargs.get('inject_solve', None) - self.function_list = [] - if level <= PERF: - self.function_list.extend([ + funcs = [ + # KSP specific 'kspgetiterationnumber', - 'snesgetiterationnumber' - ]) + 'kspgettolerances', + 'kspgetconvergedreason', + 'kspgettype', + 'kspgetnormtype', + # SNES specific + 'snesgetiterationnumber', + ] + self.query_functions = set(self.query_functions) + self.query_functions.update(funcs) + self.query_functions = sorted(list(self.query_functions)) # TODO: To be extended with if level <= DEBUG: ... @@ -32,23 +43,36 @@ def __init__(self, level, **kwargs): pname = self.sreg.make_name(prefix='petscprofiler') self.statstruct = PetscInfo( - name, pname, self.logobjs, self.sobjs, + name, pname, self.petsc_option_mapper, self.sobjs, self.section_mapper, self.inject_solve, - self.function_list + self.query_functions ) @cached_property - def logobjs(self): + def petsc_option_mapper(self): """ - Create PETSc objects specifically needed for logging solver statistics. - """ - return { - info.name: info.variable_type( - self.sreg.make_name(prefix=info.output_param) - ) - for func_name in self.function_list - for info in [petsc_return_variable_dict[func_name]] + For each function in `self.query_functions`, look up its metadata in + `petsc_return_variable_dict` and instantiate the corresponding PETSc logging + variables with names from the symbol registry. + + Example: + -------- + >>> self.query_functions + ['kspgetiterationnumber', 'snesgetiterationnumber', 'kspgettolerances'] + + >>> self.petsc_option_mapper + { + 'KSPGetIterationNumber': {'kspits': kspits0}, + 'KSPGetTolerances': {'rtol': rtol0, 'atol': atol0, ...} } + """ + opts = {} + for func_name in self.query_functions: + info = petsc_return_variable_dict[func_name] + opts[info.name] = {} + for vtype, out in zip(info.variable_type, info.output_param, strict=True): + opts[info.name][out] = vtype(self.sreg.make_name(prefix=out)) + return frozendict(opts) @cached_property def calls(self): @@ -58,20 +82,21 @@ def calls(self): """ struct = self.statstruct calls = [] - for param in self.function_list: - param = petsc_return_variable_dict[param] - - inputs = [] - for i in param.input_params: - inputs.append(self.sobjs[i]) + for func_name in self.query_functions: + return_variable = petsc_return_variable_dict[func_name] - logobj = self.logobjs[param.name] + input = self.sobjs[return_variable.input_params] + output_params = self.petsc_option_mapper[return_variable.name].values() + by_ref_output = [Byref(i) for i in output_params] calls.append( - petsc_call(param.name, inputs + [Byref(logobj)]) + petsc_call(return_variable.name, [input] + by_ref_output) ) # TODO: Perform a PetscCIntCast here? - expr = DummyExpr(FieldFromPointer(logobj._C_symbol, struct), logobj._C_symbol) - calls.append(expr) + exprs = [ + DummyExpr(FieldFromPointer(i._C_symbol, struct), i._C_symbol) + for i in output_params + ] + calls.extend(exprs) return tuple(calls) diff --git a/devito/petsc/iet/passes.py b/devito/petsc/iet/passes.py index d0118d6863..eb6cf48a1f 100644 --- a/devito/petsc/iet/passes.py +++ b/devito/petsc/iet/passes.py @@ -6,10 +6,12 @@ from devito.ir.iet import (Transformer, MapNodes, Iteration, BlankLine, DummyExpr, CallableBody, List, Call, Callable, FindNodes, Section) -from devito.symbolics import Byref, Macro, FieldFromPointer +from devito.symbolics import Byref, FieldFromPointer, Macro, Null from devito.types import Symbol, Scalar from devito.types.basic import DataSymbol from devito.tools import frozendict +import devito.logger + from devito.petsc.types import (PetscMPIInt, PetscErrorCode, MultipleFieldData, PointerIS, Mat, CallbackVec, Vec, CallbackMat, SNES, DummyArg, PetscInt, PointerDM, PointerMat, MatReuse, @@ -25,8 +27,6 @@ from devito.petsc.iet.logging import PetscLogger from devito.petsc.iet.utils import petsc_call, petsc_call_mpi -import devito.logger as dl - @iet_pass def lower_petsc(iet, **kwargs): @@ -69,9 +69,24 @@ def lower_petsc(iet, **kwargs): # Map PETScSolve to its Section (for logging) section_mapper = MapNodes(Section, PetscMetaData, 'groupby').visit(iet) + # Prefixes within the same `Operator` should not be duplicated + prefixes = [d.expr.rhs.user_prefix for d in data if d.expr.rhs.user_prefix] + duplicates = {p for p in prefixes if prefixes.count(p) > 1} + + if duplicates: + dup_list = ", ".join(repr(p) for p in sorted(duplicates)) + raise ValueError( + f"The following `options_prefix` values are duplicated " + f"among your PETScSolves. Ensure each one is unique: {dup_list}" + ) + + # List of `Call`s to clear options from the global PETSc options database, + # executed at the end of the Operator. + clear_options = [] + for iters, (inject_solve,) in inject_solve_mapper.items(): - builder = Builder(inject_solve, objs, iters, comm, section_mapper, **kwargs) + builder = Builder(inject_solve, iters, comm, section_mapper, **kwargs) setup.extend(builder.solver_setup.calls) @@ -80,11 +95,13 @@ def lower_petsc(iet, **kwargs): efuncs.update(builder.cbbuilder.efuncs) - populate_matrix_context(efuncs, objs) + clear_options.extend((petsc_call( + builder.cbbuilder._clear_options_efunc.name, [] + ),)) + populate_matrix_context(efuncs) iet = Transformer(subs).visit(iet) - - body = core + tuple(setup) + iet.body.body + body = core + tuple(setup) + iet.body.body + tuple(clear_options) body = iet.body._rebuild(body=body) iet = iet._rebuild(body=body) metadata = {**core_metadata(), 'efuncs': tuple(efuncs.values())} @@ -131,12 +148,13 @@ class Builder: returning subclasses of the objects initialised in __init__, depending on the properties of `inject_solve`. """ - def __init__(self, inject_solve, objs, iters, comm, section_mapper, **kwargs): + def __init__(self, inject_solve, iters, comm, section_mapper, **kwargs): self.inject_solve = inject_solve self.objs = objs self.iters = iters self.comm = comm self.section_mapper = section_mapper + self.get_info = inject_solve.expr.rhs.get_info self.kwargs = kwargs self.coupled = isinstance(inject_solve.expr.rhs.field_data, MultipleFieldData) self.common_kwargs = { @@ -183,15 +201,17 @@ def solve(self): @cached_property def logger(self): - log_level = dl.logger.level - return PetscLogger(log_level, **self.common_kwargs) + log_level = devito.logger.logger.level + return PetscLogger( + log_level, get_info=self.get_info, **self.common_kwargs + ) @cached_property def calls(self): return List(body=self.solve.calls+self.logger.calls) -def populate_matrix_context(efuncs, objs): +def populate_matrix_context(efuncs): if not objs['dummyefunc'] in efuncs.values(): return @@ -205,7 +225,7 @@ def populate_matrix_context(efuncs, objs): ) body = CallableBody( List(body=[subdms_expr, fields_expr]), - init=(objs['begin_user'],), + init=(petsc_func_begin_user,), retstmt=tuple([Call('PetscFunctionReturn', arguments=[0])]) ) name = 'PopulateMatContext' @@ -262,13 +282,8 @@ def populate_matrix_context(efuncs, objs): fields=[subdms, fields, submats], modifier=' *' ), 'subctx': SubMatrixStruct(fields=[rows, cols]), - 'Null': Macro('NULL'), 'dummyctx': Symbol('lctx'), 'dummyptr': DummyArg('dummy'), 'dummyefunc': Symbol('dummyefunc'), 'dof': PetscInt('dof'), - 'begin_user': c.Line('PetscFunctionBeginUser;'), }) - -# Move to macros file? -Null = Macro('NULL') diff --git a/devito/petsc/iet/routines.py b/devito/petsc/iet/routines.py index 4f7bf01e8e..94952b3842 100644 --- a/devito/petsc/iet/routines.py +++ b/devito/petsc/iet/routines.py @@ -7,7 +7,7 @@ retrieve_iteration_tree, filter_iterations, Iteration, PointerCast) from devito.symbolics import (Byref, FieldFromPointer, cast, VOID, - FieldFromComposite, IntDiv, Deref, Mod, String) + FieldFromComposite, IntDiv, Deref, Mod, String, Null) from devito.symbolics.unevaluation import Mul from devito.types.basic import AbstractFunction from devito.types import Temp, Dimension @@ -17,11 +17,11 @@ MatShellSetOp, PetscMetaData) from devito.petsc.iet.utils import (petsc_call, petsc_struct, zero_vector, dereference_funcs, residual_bundle) -from devito.petsc.utils import solver_mapper from devito.petsc.types import (PETScArray, PetscBundle, DM, Mat, CallbackVec, Vec, KSP, PC, SNES, PetscInt, StartPtr, PointerIS, PointerDM, VecScatter, DMCast, JacobianStruct, SubMatrixStruct, CallbackDM) +from devito.petsc.types.macros import petsc_func_begin_user class CBBuilder: @@ -37,10 +37,13 @@ def __init__(self, **kwargs): self.objs = kwargs.get('objs') self.solver_objs = kwargs.get('solver_objs') self.inject_solve = kwargs.get('inject_solve') + self.solve_expr = self.inject_solve.expr.rhs self._efuncs = OrderedDict() self._struct_params = [] + self._set_options_efunc = None + self._clear_options_efunc = None self._main_matvec_callback = None self._user_struct_callback = None self._F_efunc = None @@ -92,9 +95,17 @@ def initial_guesses(self): def user_struct_callback(self): return self._user_struct_callback + @property + def solver_parameters(self): + return self.solve_expr.solver_parameters + @property def field_data(self): - return self.inject_solve.expr.rhs.field_data + return self.solve_expr.field_data + + @property + def formatted_prefix(self): + return self.solve_expr.formatted_prefix @property def arrays(self): @@ -105,6 +116,7 @@ def target(self): return self.field_data.target def _make_core(self): + self._make_options_callback() self._make_matvec(self.field_data.jacobian) self._make_formfunc() self._make_formrhs() @@ -112,6 +124,69 @@ def _make_core(self): self._make_initial_guess() self._make_user_struct_callback() + def _make_petsc_callable(self, prefix, body, parameters=()): + return PETScCallable( + self.sregistry.make_name(prefix=prefix), + body, + retval=self.objs['err'], + parameters=parameters + ) + + def _make_callable_body(self, body, stacks=(), casts=()): + return CallableBody( + List(body=body), + init=(petsc_func_begin_user,), + stacks=stacks, + casts=casts, + retstmt=(Call('PetscFunctionReturn', arguments=[0]),) + ) + + def _make_options_callback(self): + """ + Create two callbacks: one to set PETSc options and one + to clear them. + + Options are only set/cleared if they were not specifed via + command line arguments. + """ + params = self.solver_parameters + prefix = self.formatted_prefix + + set_body, clear_body = [], [] + + for k, v in params.items(): + option = f'-{prefix}{k}' + + # TODO: Revisit use of a global variable here. + # Consider replacing this with a call to `PetscGetArgs`, though + # initial attempts failed, possibly because the argv pointer is + # created in Python?.. + import devito.petsc.initialize + if option in devito.petsc.initialize._petsc_clargs: + # Ensures that the command line args take priority + continue + + option_name = String(option) + # For options without a value e.g `ksp_view`, pass Null + option_value = Null if v is None else String(str(v)) + set_body.append( + petsc_call('PetscOptionsSetValue', [Null, option_name, option_value]) + ) + clear_body.append( + petsc_call('PetscOptionsClearValue', [Null, option_name]) + ) + + set_body = self._make_callable_body(set_body) + clear_body = self._make_callable_body(clear_body) + + set_callback = self._make_petsc_callable('SetPetscOptions', set_body) + clear_callback = self._make_petsc_callable('ClearPetscOptions', clear_body) + + self._set_options_efunc = set_callback + self._efuncs[set_callback.name] = set_callback + self._clear_options_efunc = clear_callback + self._efuncs[clear_callback.name] = clear_callback + def _make_matvec(self, jacobian, prefix='MatMult'): # Compile `matvecs` into an IET via recursive compilation matvecs = jacobian.matvecs @@ -122,13 +197,9 @@ def _make_matvec(self, jacobian, prefix='MatMult'): body = self._create_matvec_body( List(body=irs.uiet.body), jacobian ) - objs = self.objs - cb = PETScCallable( - self.sregistry.make_name(prefix=prefix), - body, - retval=objs['err'], - parameters=(objs['J'], objs['X'], objs['Y']) + cb = self._make_petsc_callable( + prefix, body, parameters=(objs['J'], objs['X'], objs['Y']) ) self._J_efuncs.append(cb) self._efuncs[cb.name] = cb @@ -246,12 +317,7 @@ def _create_matvec_body(self, body, jacobian): # Dereference function data in struct derefs = dereference_funcs(ctx, fields) - body = CallableBody( - List(body=body), - init=(objs['begin_user'],), - stacks=stacks+derefs, - retstmt=(Call('PetscFunctionReturn', arguments=[0]),) - ) + body = self._make_callable_body(body, stacks=stacks+derefs) # Replace non-function data with pointer to data in struct subs = {i._C_symbol: FieldFromPointer(i._C_symbol, ctx) for i in fields} @@ -261,6 +327,7 @@ def _create_matvec_body(self, body, jacobian): return body def _make_formfunc(self): + objs = self.objs F_exprs = self.field_data.residual.F_exprs # Compile `F_exprs` into an IET via recursive compilation irs, _ = self.rcompile( @@ -270,13 +337,9 @@ def _make_formfunc(self): body_formfunc = self._create_formfunc_body( List(body=irs.uiet.body) ) - objs = self.objs - cb = PETScCallable( - self.sregistry.make_name(prefix='FormFunction'), - body_formfunc, - retval=objs['err'], - parameters=(objs['snes'], objs['X'], objs['F'], objs['dummyptr']) - ) + parameters = (objs['snes'], objs['X'], objs['F'], objs['dummyptr']) + cb = self._make_petsc_callable('FormFunction', body_formfunc, parameters) + self._F_efunc = cb self._efuncs[cb.name] = cb @@ -385,12 +448,7 @@ def _create_formfunc_body(self, body): # Dereference function data in struct derefs = dereference_funcs(ctx, fields) - body = CallableBody( - List(body=body), - init=(objs['begin_user'],), - stacks=stacks+derefs, - retstmt=(Call('PetscFunctionReturn', arguments=[0]),)) - + body = self._make_callable_body(body, stacks=stacks+derefs) # Replace non-function data with pointer to data in struct subs = {i._C_symbol: FieldFromPointer(i._C_symbol, ctx) for i in fields} @@ -409,11 +467,8 @@ def _make_formrhs(self): List(body=irs.uiet.body) ) objs = self.objs - cb = PETScCallable( - self.sregistry.make_name(prefix='FormRHS'), - body, - retval=objs['err'], - parameters=(sobjs['callbackdm'], objs['B']) + cb = self._make_petsc_callable( + 'FormRHS', body, parameters=(sobjs['callbackdm'], objs['B']) ) self._b_efunc = cb self._efuncs[cb.name] = cb @@ -495,12 +550,7 @@ def _create_form_rhs_body(self, body): # Dereference function data in struct derefs = dereference_funcs(ctx, fields) - body = CallableBody( - List(body=[body]), - init=(objs['begin_user'],), - stacks=stacks+derefs, - retstmt=(Call('PetscFunctionReturn', arguments=[0]),) - ) + body = self._make_callable_body([body], stacks=stacks+derefs) # Replace non-function data with pointer to data in struct subs = {i._C_symbol: FieldFromPointer(i._C_symbol, ctx) for @@ -511,6 +561,7 @@ def _create_form_rhs_body(self, body): def _make_initial_guess(self): exprs = self.field_data.initial_guess.exprs sobjs = self.solver_objs + objs = self.objs # Compile initital guess `eqns` into an IET via recursive compilation irs, _ = self.rcompile( @@ -520,12 +571,8 @@ def _make_initial_guess(self): body = self._create_initial_guess_body( List(body=irs.uiet.body) ) - objs = self.objs - cb = PETScCallable( - self.sregistry.make_name(prefix='FormInitialGuess'), - body, - retval=objs['err'], - parameters=(sobjs['callbackdm'], objs['xloc']) + cb = self._make_petsc_callable( + 'FormInitialGuess', body, parameters=(sobjs['callbackdm'], objs['xloc']) ) self._initial_guesses.append(cb) self._efuncs[cb.name] = cb @@ -572,13 +619,7 @@ def _create_initial_guess_body(self, body): # Dereference function data in struct derefs = dereference_funcs(ctx, fields) - - body = CallableBody( - List(body=[body]), - init=(objs['begin_user'],), - stacks=stacks+derefs, - retstmt=(Call('PetscFunctionReturn', arguments=[0]),) - ) + body = self._make_callable_body(body, stacks=stacks+derefs) # Replace non-function data with pointer to data in struct subs = {i._C_symbol: FieldFromPointer(i._C_symbol, ctx) for @@ -601,10 +642,7 @@ def _make_user_struct_callback(self): DummyExpr(FieldFromPointer(i._C_symbol, mainctx), i._C_symbol) for i in mainctx.callback_fields ] - struct_callback_body = CallableBody( - List(body=body), init=(self.objs['begin_user'],), - retstmt=(Call('PetscFunctionReturn', arguments=[0]),) - ) + struct_callback_body = self._make_callable_body(body) cb = Callable( self.sregistry.make_name(prefix='PopulateUserContext'), struct_callback_body, self.objs['err'], @@ -663,6 +701,7 @@ def _make_core(self): for sm in self.field_data.jacobian.nonzero_submatrices: self._make_matvec(sm, prefix=f'{sm.name}_MatMult') + self._make_options_callback() self._make_whole_matvec() self._make_whole_formfunc() self._make_user_struct_callback() @@ -673,11 +712,9 @@ def _make_whole_matvec(self): objs = self.objs body = self._whole_matvec_body() - cb = PETScCallable( - self.sregistry.make_name(prefix='WholeMatMult'), - List(body=body), - retval=objs['err'], - parameters=(objs['J'], objs['X'], objs['Y']) + parameters = (objs['J'], objs['X'], objs['Y']) + cb = self._make_petsc_callable( + 'WholeMatMult', List(body=body), parameters=parameters ) self._main_matvec_callback = cb self._efuncs[cb.name] = cb @@ -724,13 +761,11 @@ def _whole_matvec_body(self): [objs['Y'], Deref(FieldFromPointer(rows, ctx)), Byref(Y)] ), ) - return CallableBody( - List(body=(ctx_main, zero_y_memory, BlankLine) + calls), - init=(objs['begin_user'],), - retstmt=(Call('PetscFunctionReturn', arguments=[0]),) - ) + body = (ctx_main, zero_y_memory, BlankLine) + calls + return self._make_callable_body(body) def _make_whole_formfunc(self): + objs = self.objs F_exprs = self.field_data.residual.F_exprs # Compile `F_exprs` into an IET via recursive compilation irs, _ = self.rcompile( @@ -739,13 +774,11 @@ def _make_whole_formfunc(self): ) body = self._whole_formfunc_body(List(body=irs.uiet.body)) - objs = self.objs - cb = PETScCallable( - self.sregistry.make_name(prefix='WholeFormFunc'), - body, - retval=objs['err'], - parameters=(objs['snes'], objs['X'], objs['F'], objs['dummyptr']) + parameters = (objs['snes'], objs['X'], objs['F'], objs['dummyptr']) + cb = self._make_petsc_callable( + 'WholeFormFunc', body, parameters=parameters ) + self._F_efunc = cb self._efuncs[cb.name] = cb @@ -859,14 +892,10 @@ def _whole_formfunc_body(self, body): f_soa = PointerCast(fbundle) x_soa = PointerCast(xbundle) - formfunc_body = CallableBody( - List(body=body), - init=(objs['begin_user'],), - stacks=stacks+derefs, + formfunc_body = self._make_callable_body( + body, stacks=stacks+derefs, casts=(f_soa, x_soa), - retstmt=(Call('PetscFunctionReturn', arguments=[0]),) ) - # Replace non-function data with pointer to data in struct subs = {i._C_symbol: FieldFromPointer(i._C_symbol, ctx) for i in fields} @@ -883,12 +912,9 @@ def _create_submatrices(self): objs['matreuse'], objs['Submats'], ) - cb = PETScCallable( - self.sregistry.make_name(prefix='MatCreateSubMatrices'), - body, - retval=objs['err'], - parameters=params - ) + cb = self._make_petsc_callable( + 'MatCreateSubMatrices', body, parameters=params) + self._submatrices_callback = cb self._efuncs[cb.name] = cb @@ -910,7 +936,6 @@ def _submat_callback_body(self): get_ctx = petsc_call('MatShellGetContext', [objs['J'], Byref(objs['ljacctx'])]) - Null = objs['Null'] dm_get_info = petsc_call( 'DMDAGetInfo', [ sobjs['callbackdm'], Null, Byref(sobjs['M']), Byref(sobjs['N']), @@ -1011,12 +1036,7 @@ def _submat_callback_body(self): iteration, ] + matmult_op - return CallableBody( - List(body=tuple(body)), - init=(objs['begin_user'],), - stacks=(get_ctx, deref_subdm), - retstmt=(Call('PetscFunctionReturn', arguments=[0]),) - ) + return self._make_callable_body(tuple(body), stacks=(get_ctx, deref_subdm)) class BaseObjectBuilder: @@ -1058,7 +1078,7 @@ def _build(self): targets = self.field_data.targets snes_name = sreg.make_name(prefix='snes') - options_prefix = self.inject_solve.expr.rhs.options_prefix + formatted_prefix = self.inject_solve.expr.rhs.formatted_prefix base_dict = { 'Jac': Mat(sreg.make_name(prefix='J')), @@ -1072,9 +1092,9 @@ def _build(self): 'localsize': PetscInt(sreg.make_name(prefix='localsize')), 'dmda': DM(sreg.make_name(prefix='da'), dofs=len(targets)), 'callbackdm': CallbackDM(sreg.make_name(prefix='dm')), - 'snesprefix': String((options_prefix or '') + '_'), - 'options_prefix': options_prefix, + 'snes_prefix': String(formatted_prefix), } + base_dict['comm'] = self.comm self._target_dependent(base_dict) return self._extend_build(base_dict) @@ -1213,6 +1233,7 @@ def __init__(self, **kwargs): self.solver_objs = kwargs.get('solver_objs') self.cbbuilder = kwargs.get('cbbuilder') self.field_data = self.inject_solve.expr.rhs.field_data + self.formatted_prefix = self.inject_solve.expr.rhs.formatted_prefix self.calls = self._setup() @property @@ -1224,29 +1245,26 @@ def snes_ctx(self): return VOID(self.solver_objs['dmda'], stars='*') def _setup(self): - objs = self.objs sobjs = self.solver_objs - dmda = sobjs['dmda'] - solver_params = self.inject_solve.expr.rhs.solver_parameters - snes_create = petsc_call('SNESCreate', [sobjs['comm'], Byref(sobjs['snes'])]) snes_options_prefix = petsc_call( - 'SNESSetOptionsPrefix', [sobjs['snes'], sobjs['snesprefix']] - ) if sobjs['options_prefix'] else None + 'SNESSetOptionsPrefix', [sobjs['snes'], sobjs['snes_prefix']] + ) if self.formatted_prefix else None + + set_options = petsc_call( + self.cbbuilder._set_options_efunc.name, [] + ) snes_set_dm = petsc_call('SNESSetDM', [sobjs['snes'], dmda]) create_matrix = petsc_call('DMCreateMatrix', [dmda, Byref(sobjs['Jac'])]) - # NOTE: Assuming all solves are linear for now - snes_set_type = petsc_call('SNESSetType', [sobjs['snes'], 'SNESKSPONLY']) - snes_set_jac = petsc_call( 'SNESSetJacobian', [sobjs['snes'], sobjs['Jac'], - sobjs['Jac'], 'MatMFFDComputeJacobian', objs['Null']] + sobjs['Jac'], 'MatMFFDComputeJacobian', Null] ) global_x = petsc_call('DMCreateGlobalVector', @@ -1260,6 +1278,7 @@ def _setup(self): local_size = math.prod( v for v, dim in zip(target.shape_allocated, target.dimensions) if dim.is_Space ) + # TODO: Check - VecCreateSeqWithArray local_x = petsc_call('VecCreateMPIWithArray', [sobjs['comm'], 1, local_size, 'PETSC_DECIDE', field_from_ptr, Byref(sobjs['xlocal'])]) @@ -1275,25 +1294,6 @@ def _setup(self): snes_get_ksp = petsc_call('SNESGetKSP', [sobjs['snes'], Byref(sobjs['ksp'])]) - ksp_set_tols = petsc_call( - 'KSPSetTolerances', [sobjs['ksp'], solver_params['ksp_rtol'], - solver_params['ksp_atol'], solver_params['ksp_divtol'], - solver_params['ksp_max_it']] - ) - - ksp_set_type = petsc_call( - 'KSPSetType', [sobjs['ksp'], solver_mapper[solver_params['ksp_type']]] - ) - - ksp_get_pc = petsc_call( - 'KSPGetPC', [sobjs['ksp'], Byref(sobjs['pc'])] - ) - - # Even though the default will be jacobi, set to PCNONE for now - pc_set_type = petsc_call('PCSetType', [sobjs['pc'], 'PCNONE']) - - ksp_set_from_ops = petsc_call('KSPSetFromOptions', [sobjs['ksp']]) - matvec = self.cbbuilder.main_matvec_callback matvec_operation = petsc_call( 'MatShellSetOperation', @@ -1302,7 +1302,7 @@ def _setup(self): formfunc = self.cbbuilder._F_efunc formfunc_operation = petsc_call( 'SNESSetFunction', - [sobjs['snes'], objs['Null'], FormFunctionCallback(formfunc.name, void, void), + [sobjs['snes'], Null, FormFunctionCallback(formfunc.name, void, void), self.snes_ctx] ) @@ -1326,20 +1326,15 @@ def _setup(self): base_setup = dmda_calls + ( snes_create, snes_options_prefix, + set_options, snes_set_dm, create_matrix, snes_set_jac, - snes_set_type, global_x, local_x, get_local_size, global_b, snes_get_ksp, - ksp_set_tols, - ksp_set_type, - ksp_get_pc, - pc_set_type, - ksp_set_from_ops, matvec_operation, formfunc_operation, snes_set_options, @@ -1364,7 +1359,6 @@ def _create_dmda_calls(self, dmda): return dmda_create, dm_setup, dm_mat_type def _create_dmda(self, dmda): - objs = self.objs sobjs = self.solver_objs grid = self.field_data.grid nspace_dims = len(grid.dimensions) @@ -1393,7 +1387,7 @@ def _create_dmda(self, dmda): stencil_width = self.field_data.space_order args.append(stencil_width) - args.extend([objs['Null']]*nspace_dims) + args.extend([Null]*nspace_dims) # The distributed array object args.append(Byref(dmda)) @@ -1409,26 +1403,25 @@ def _setup(self): # TODO: minimise code duplication with superclass objs = self.objs sobjs = self.solver_objs - dmda = sobjs['dmda'] - solver_params = self.inject_solve.expr.rhs.solver_parameters snes_create = petsc_call('SNESCreate', [sobjs['comm'], Byref(sobjs['snes'])]) snes_options_prefix = petsc_call( - 'SNESSetOptionsPrefix', [sobjs['snes'], sobjs['snesprefix']] - ) if sobjs['options_prefix'] else None + 'SNESSetOptionsPrefix', [sobjs['snes'], sobjs['snes_prefix']] + ) if self.formatted_prefix else None + + set_options = petsc_call( + self.cbbuilder._set_options_efunc.name, [] + ) snes_set_dm = petsc_call('SNESSetDM', [sobjs['snes'], dmda]) create_matrix = petsc_call('DMCreateMatrix', [dmda, Byref(sobjs['Jac'])]) - # NOTE: Assuming all solves are linear for now - snes_set_type = petsc_call('SNESSetType', [sobjs['snes'], 'SNESKSPONLY']) - snes_set_jac = petsc_call( 'SNESSetJacobian', [sobjs['snes'], sobjs['Jac'], - sobjs['Jac'], 'MatMFFDComputeJacobian', objs['Null']] + sobjs['Jac'], 'MatMFFDComputeJacobian', Null] ) global_x = petsc_call('DMCreateGlobalVector', @@ -1442,25 +1435,6 @@ def _setup(self): snes_get_ksp = petsc_call('SNESGetKSP', [sobjs['snes'], Byref(sobjs['ksp'])]) - ksp_set_tols = petsc_call( - 'KSPSetTolerances', [sobjs['ksp'], solver_params['ksp_rtol'], - solver_params['ksp_atol'], solver_params['ksp_divtol'], - solver_params['ksp_max_it']] - ) - - ksp_set_type = petsc_call( - 'KSPSetType', [sobjs['ksp'], solver_mapper[solver_params['ksp_type']]] - ) - - ksp_get_pc = petsc_call( - 'KSPGetPC', [sobjs['ksp'], Byref(sobjs['pc'])] - ) - - # Even though the default will be jacobi, set to PCNONE for now - pc_set_type = petsc_call('PCSetType', [sobjs['pc'], 'PCNONE']) - - ksp_set_from_ops = petsc_call('KSPSetFromOptions', [sobjs['ksp']]) - matvec = self.cbbuilder.main_matvec_callback matvec_operation = petsc_call( 'MatShellSetOperation', @@ -1469,7 +1443,7 @@ def _setup(self): formfunc = self.cbbuilder._F_efunc formfunc_operation = petsc_call( 'SNESSetFunction', - [sobjs['snes'], objs['Null'], FormFunctionCallback(formfunc.name, void, void), + [sobjs['snes'], Null, FormFunctionCallback(formfunc.name, void, void), self.snes_ctx] ) @@ -1492,7 +1466,7 @@ def _setup(self): create_field_decomp = petsc_call( 'DMCreateFieldDecomposition', - [dmda, Byref(sobjs['nfields']), objs['Null'], Byref(sobjs['fields']), + [dmda, Byref(sobjs['nfields']), Null, Byref(sobjs['fields']), Byref(sobjs['subdms'])] ) submat_cb = self.cbbuilder.submatrices_callback @@ -1533,19 +1507,14 @@ def _setup(self): coupled_setup = dmda_calls + ( snes_create, snes_options_prefix, + set_options, snes_set_dm, create_matrix, snes_set_jac, - snes_set_type, global_x, local_x, get_local_size, snes_get_ksp, - ksp_set_tols, - ksp_set_type, - ksp_get_pc, - pc_set_type, - ksp_set_from_ops, matvec_operation, formfunc_operation, snes_set_options, @@ -1642,7 +1611,6 @@ def _execute_solve(self): Assigns the required time iterators to the struct and executes the necessary calls to execute the SNES solver. """ - objs = self.objs sobjs = self.solver_objs xglob = sobjs['xglobal'] @@ -1673,7 +1641,7 @@ def _execute_solve(self): ), petsc_call( 'VecScatterCreate', - [xglob, field, target_xglob, self.objs['Null'], Byref(s)] + [xglob, field, target_xglob, Null, Byref(s)] ), petsc_call( 'VecScatterBegin', @@ -1701,7 +1669,7 @@ def _execute_solve(self): ) ) - snes_solve = (petsc_call('SNESSolve', [sobjs['snes'], objs['Null'], xglob]),) + snes_solve = (petsc_call('SNESSolve', [sobjs['snes'], Null, xglob]),) return (struct_assignment,) + pre_solve + snes_solve + post_solve + (BlankLine,) @@ -1758,10 +1726,10 @@ class TimeDependent(NonTimeDependent): for each `SNESSolve` at every time step, don't require the time loop, but may still need access to data from other time steps. - All `Function` objects are passed through the initial lowering via the - `LinearSolveExpr` object, ensuring the correct time loop is generated + `SolverMetaData` object, ensuring the correct time loop is generated in the main kernel. - Another mapper is created based on the modulo dimensions - generated by the `LinearSolveExpr` object in the main kernel + generated by the `SolverMetaData` object in the main kernel (e.g., {time: time, t: t0, t + 1: t1}). - These two mappers are used to generate a final mapper `symb_to_moddim` (e.g. {tau0: t0, tau1: t1}) which is used at the IET level to diff --git a/devito/petsc/iet/utils.py b/devito/petsc/iet/utils.py index 3402cfe28f..d7ccfb2b4e 100644 --- a/devito/petsc/iet/utils.py +++ b/devito/petsc/iet/utils.py @@ -1,8 +1,9 @@ -from devito.petsc.iet.nodes import PetscMetaData, PETScCall from devito.ir.equations import OpPetsc from devito.ir.iet import Dereference, FindSymbols, Uxreplace from devito.types.basic import AbstractFunction +from devito.petsc.iet.nodes import PetscMetaData, PETScCall + def petsc_call(specific_call, call_args): return PETScCall('PetscCall', [PETScCall(specific_call, arguments=call_args)]) diff --git a/devito/petsc/initialize.py b/devito/petsc/initialize.py index 80e3c7520c..a4c136f71a 100644 --- a/devito/petsc/initialize.py +++ b/devito/petsc/initialize.py @@ -12,10 +12,21 @@ _petsc_initialized = False -def PetscInitialize(): +global _petsc_clargs + + +def PetscInitialize(clargs=sys.argv): global _petsc_initialized + global _petsc_clargs + if not _petsc_initialized: dummy = Symbol(name='d') + + if clargs is not sys.argv: + clargs = (sys.argv[0], *clargs) + + _petsc_clargs = clargs + # TODO: Potentially just use cgen + the compiler machinery in Devito # to generate these "dummy_ops" instead of using the Operator class. # This would prevent circular imports when initializing during import @@ -30,14 +41,19 @@ def PetscInitialize(): name='kernel_finalize', opt='noop' ) + # Convert each string to a bytes object (e.g: '-ksp_type' -> b'-ksp_type') # `argv_bytes` must be a list so the memory address persists # `os.fsencode` should be preferred over `string().encode('utf-8')` # in case there is some system specific encoding in use - argv_bytes = list(map(os.fsencode, sys.argv)) - argv_pointer = (POINTER(c_char)*len(sys.argv))( + argv_bytes = list(map(os.fsencode, clargs)) + + # POINTER(c_char) is equivalent to char * in C + # (POINTER(c_char) * len(clargs)) creates a C array type: char *[len(clargs)] + # Instantiating it with (*map(...)) casts each bytes object to a char * and + # fills the array. The result is a char *argv[] + argv_pointer = (POINTER(c_char)*len(clargs))( *map(lambda s: cast(s, POINTER(c_char)), argv_bytes) ) - op_init.apply(argc=len(sys.argv), argv=argv_pointer) - + op_init.apply(argc=len(clargs), argv=argv_pointer) atexit.register(op_finalize.apply) _petsc_initialized = True diff --git a/devito/petsc/logging.py b/devito/petsc/logging.py index 40ad2711fa..979b4b582a 100644 --- a/devito/petsc/logging.py +++ b/devito/petsc/logging.py @@ -1,10 +1,13 @@ -from collections import namedtuple, OrderedDict +import os +from collections import namedtuple from dataclasses import dataclass from devito.types import CompositeObject -from devito.petsc.types import PetscInt -from devito.petsc.utils import petsc_type_mappings +from devito.petsc.types import ( + PetscInt, PetscScalar, KSPType, KSPConvergedReason, KSPNormType +) +from devito.petsc.utils import petsc_type_to_ctype class PetscEntry: @@ -17,6 +20,9 @@ def __init__(self, **kwargs): def __getitem__(self, key): return self._properties[key.lower()] + def __len__(self): + return len(self._properties) + def __repr__(self): return f"PetscEntry({', '.join(f'{k}={v}' for k, v in self.kwargs.items())})" @@ -37,7 +43,7 @@ def __init__(self, params, *args, **kwargs): self._functions = list(dict.fromkeys( petsc_return_variable_dict[key].name for struct in self.petscinfos - for key in struct.function_list + for key in struct.query_functions )) self._property_name_map = {} # Dynamically create a property on this class for each PETSc function @@ -62,8 +68,12 @@ def petsc_entry(self, petscinfo): Create a named tuple entry for the given PetscInfo object, containing the values for each PETSc function call. """ - funcs = self._functions - values = tuple(getattr(petscinfo, c) for c in funcs) + # Collect the function names associated with this PetscInfo + # instance (i.e., for a single PETScSolve). + funcs = [ + petsc_return_variable_dict[f].name for f in petscinfo.query_functions + ] + values = [getattr(petscinfo, c) for c in funcs] return PetscEntry(**{k: v for k, v in zip(funcs, values)}) def _add_properties(self): @@ -71,18 +81,18 @@ def _add_properties(self): For each function name in `self._functions` (e.g., 'KSPGetIterationNumber'), dynamically add a property to the class with the same name. - Each property returns an OrderedDict that maps each PetscKey to the + Each property returns a dict mapping each PetscKey to the result of looking up that function on the corresponding PetscEntry, if the function exists on that entry. """ def make_property(function): def getter(self): - return OrderedDict( - (k, getattr(v, function)) + return { + k: getattr(v, function) for k, v in self.items() # Only include entries that have the function if hasattr(v, function) - ) + } return property(getter) for f in self._functions: @@ -119,21 +129,42 @@ def __getitem__(self, key): class PetscInfo(CompositeObject): - __rargs__ = ('name', 'pname', 'logobjs', 'sobjs', 'section_mapper', - 'inject_solve', 'function_list') + __rargs__ = ('name', 'pname', 'petsc_option_mapper', 'sobjs', 'section_mapper', + 'inject_solve', 'query_functions') - def __init__(self, name, pname, logobjs, sobjs, section_mapper, - inject_solve, function_list): + def __init__(self, name, pname, petsc_option_mapper, sobjs, section_mapper, + inject_solve, query_functions): - self.logobjs = logobjs + self.petsc_option_mapper = petsc_option_mapper self.sobjs = sobjs self.section_mapper = section_mapper self.inject_solve = inject_solve - self.function_list = function_list + self.query_functions = query_functions + self.solve_expr = inject_solve.expr.rhs + + pfields = [] + + # All petsc options needed to form the PetscInfo struct + # e.g (kspits0, rtol0, atol0, ...) + self._fields = [i for j in petsc_option_mapper.values() for i in j.values()] + + for petsc_option in self._fields: + # petsc_type is e.g. 'PetscInt', 'PetscScalar', 'KSPType' + petsc_type = str(petsc_option.dtype) + ctype = petsc_type_to_ctype[petsc_type] + pfields.append((petsc_option.name, ctype)) - mapper = {v: k for k, v in petsc_type_mappings.items()} - fields = [(str(i), mapper[str(i._C_ctype)]) for i in logobjs.values()] - super().__init__(name, pname, fields) + super().__init__(name, pname, pfields) + + @property + def fields(self): + return self._fields + + @property + def prefix(self): + # If users provide an options prefix, use it in the summary; + # otherwise, use the default one generated by Devito + return self.solve_expr.user_prefix or self.solve_expr.formatted_prefix @property def section(self): @@ -142,20 +173,38 @@ def section(self): @property def summary_key(self): - return (self.section, self.sobjs['options_prefix']) + return (self.section, self.prefix) def __getattr__(self, attr): - if attr in self.logobjs.keys(): - return getattr(self.value._obj, self.logobjs[attr].name) - raise AttributeError(f"{attr} not found in PETSc return variables") + if attr not in self.petsc_option_mapper: + raise AttributeError(f"{attr} not found in PETSc return variables") + + # Maps the petsc_option to its generated variable name e.g {'its': its0} + obj_mapper = self.petsc_option_mapper[attr] + + def get_val(val): + val = getattr(self.value._obj, val.name) + # Decode the val if it is a bytes object + return str(os.fsdecode(val)) if isinstance(val, bytes) else val + + # - If the function returns a single value (e.g., KSPGetIterationNumber), + # return that value directly. + # - If the function returns multiple values (e.g., KSPGetTolerances), + # return a dictionary mapping each output name to its value, + # e.g., {'rtol': val0, 'atol': val1, ...}. + info = {k: get_val(v) for k, v in obj_mapper.items()} + if len(info) == 1: + return info.popitem()[1] + else: + return info @dataclass class PetscReturnVariable: name: str - variable_type: None - input_params: list - output_param: str + variable_type: tuple + input_params: str + output_param: tuple[str] # NOTE: @@ -164,17 +213,44 @@ class PetscReturnVariable: # If any of the PETSc function signatures change (e.g., names, input/output parameters), # this dictionary must be updated accordingly. +# TODO: To be extended petsc_return_variable_dict = { + # KSP specific 'kspgetiterationnumber': PetscReturnVariable( name='KSPGetIterationNumber', - variable_type=PetscInt, - input_params=['ksp'], - output_param='kspiter' + variable_type=(PetscInt,), + input_params='ksp', + output_param=('kspits',) + ), + 'kspgettolerances': PetscReturnVariable( + name='KSPGetTolerances', + variable_type=(PetscScalar, PetscScalar, PetscScalar, PetscInt), + input_params='ksp', + output_param=('rtol', 'atol', 'divtol', 'max_it'), ), + 'kspgetconvergedreason': PetscReturnVariable( + name='KSPGetConvergedReason', + variable_type=(KSPConvergedReason,), + input_params='ksp', + output_param=('reason',), + ), + 'kspgettype': PetscReturnVariable( + name='KSPGetType', + variable_type=(KSPType,), + input_params='ksp', + output_param=('ksptype',), + ), + 'kspgetnormtype': PetscReturnVariable( + name='KSPGetNormType', + variable_type=(KSPNormType,), + input_params='ksp', + output_param=('kspnormtype',), + ), + # SNES specific -> will be extended when non-linear solvers are supported 'snesgetiterationnumber': PetscReturnVariable( name='SNESGetIterationNumber', - variable_type=PetscInt, - input_params=['snes'], - output_param='snesiter', - ) + variable_type=(PetscInt,), + input_params='snes', + output_param=('snesits',), + ), } diff --git a/devito/petsc/solve.py b/devito/petsc/solve.py index 77022fc6b8..86b6d624e4 100644 --- a/devito/petsc/solve.py +++ b/devito/petsc/solve.py @@ -1,18 +1,20 @@ -from devito.types import Symbol, SteppingDimension from devito.types.equation import PetscEq -from devito.operations.solve import eval_time_derivatives -from devito.symbolics import retrieve_functions -from devito.tools import as_tuple, filter_ordered -from devito.petsc.types import (LinearSolveExpr, PETScArray, DMDALocalInfo, +from devito.tools import as_tuple +from devito.petsc.types import (LinearSolverMetaData, PETScArray, DMDALocalInfo, FieldData, MultipleFieldData, Jacobian, Residual, MixedResidual, MixedJacobian, InitialGuess) from devito.petsc.types.equation import EssentialBC +from devito.petsc.solver_parameters import (linear_solver_parameters, + format_options_prefix) +from devito.petsc.utils import get_funcs, generate_time_mapper __all__ = ['PETScSolve'] -def PETScSolve(target_exprs, target=None, solver_parameters=None, options_prefix=None): +# TODO: Rename this to petsc_solve, petscsolve? +def PETScSolve(target_exprs, target=None, solver_parameters=None, + options_prefix=None, get_info=[]): """ Returns a symbolic expression representing a linear PETSc solver, enriched with all the necessary metadata for execution within an `Operator`. @@ -44,6 +46,34 @@ def PETScSolve(target_exprs, target=None, solver_parameters=None, options_prefix solver_parameters : dict, optional PETSc solver options. + Descriptions (not exhaustive): + - 'ksp_type': Specifies the Krylov method (e.g., 'gmres', 'cg'). + - 'pc_type': Specifies the preconditioner type (e.g., ...). + - 'ksp_rtol': Relative convergence tolerance for KSP solvers. + - 'ksp_atol': Absolute convergence tolerance for KSP solvers. + - 'ksp_divtol': Divergence tolerance, amount residual norm can increase before + `KSPConvergedDefault()` concludes that the method is diverging. + - 'ksp_max_it': Maximum number of KSP iterations to use. + - 'snes_type': Type of SNES solver; 'ksponly' is used for linear solves. + + References: + - KSP types: https://petsc.org/main/manualpages/KSP/KSPType/ + - PC types: https://petsc.org/main/manualpages/PC/PCType/ + - KSP tolerances: https://petsc.org/main/manualpages/KSP/KSPSetTolerances/ + - SNES type: https://petsc.org/main/manualpages/SNES/SNESType/ + + options_prefix : str, optional + Prefix for the solver, used to configure options via the command line. If not + provided, a default prefix is generated by Devito. + + get_info : list[str], optional + A list of PETSc API functions (case insensitive) to collect statistics + from the solver. + + List of available functions: + - ['kspgetiterationnumber', 'kspgettolerances', 'kspgetconvergedreason', + 'kspgettype', 'kspgetnormtype', 'snesgetiterationnumber'] + Returns ------- Eq: @@ -51,33 +81,37 @@ def PETScSolve(target_exprs, target=None, solver_parameters=None, options_prefix This can be passed directly to a Devito Operator. """ if target is not None: - return InjectSolve( - solver_parameters, {target: target_exprs}, options_prefix - ).build_expr() + return InjectSolve(solver_parameters, {target: target_exprs}, + options_prefix, get_info).build_expr() else: - return InjectMixedSolve( - solver_parameters, target_exprs, options_prefix - ).build_expr() + return InjectMixedSolve(solver_parameters, target_exprs, + options_prefix, get_info).build_expr() class InjectSolve: - def __init__(self, solver_parameters=None, target_exprs=None, options_prefix=None): - self.solver_params = solver_parameters + def __init__(self, solver_parameters=None, target_exprs=None, options_prefix=None, + get_info=[]): + self.solver_parameters = linear_solver_parameters(solver_parameters) self.time_mapper = None self.target_exprs = target_exprs - self.options_prefix = options_prefix + # The original options prefix provided by the user + self.user_prefix = options_prefix + self.formatted_prefix = format_options_prefix(options_prefix) + self.get_info = [f.lower() for f in get_info] def build_expr(self): target, funcs, field_data = self.linear_solve_args() - # Placeholder expression for inserting calls to the solver - linear_solve = LinearSolveExpr( + + linear_solve = LinearSolverMetaData( funcs, - self.solver_params, + solver_parameters=self.solver_parameters, field_data=field_data, time_mapper=self.time_mapper, localinfo=localinfo, - options_prefix=self.options_prefix + user_prefix=self.user_prefix, + formatted_prefix=self.formatted_prefix, + get_info=self.get_info ) return PetscEq(target, linear_solve) @@ -152,45 +186,5 @@ def linear_solve_args(self): return targets[0], tuple(funcs), all_data -def generate_time_mapper(funcs): - """ - Replace time indices with `Symbols` in expressions used within - PETSc callback functions. These symbols are Uxreplaced at the IET - level to align with the `TimeDimension` and `ModuloDimension` objects - present in the initial lowering. - NOTE: All functions used in PETSc callback functions are attached to - the `LinearSolveExpr` object, which is passed through the initial lowering - (and subsequently dropped and replaced with calls to run the solver). - Therefore, the appropriate time loop will always be correctly generated inside - the main kernel. - Examples - -------- - >>> funcs = [ - >>> f1(t + dt, x, y), - >>> g1(t + dt, x, y), - >>> g2(t, x, y), - >>> f1(t, x, y) - >>> ] - >>> generate_time_mapper(funcs) - {t + dt: tau0, t: tau1} - """ - time_indices = list({ - i if isinstance(d, SteppingDimension) else d - 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 dict(zip(time_indices, tau_symbs)) - - -def get_funcs(exprs): - funcs = [ - f for e in exprs - for f in retrieve_functions(eval_time_derivatives(e.lhs - e.rhs)) - ] - return filter_ordered(funcs) - - localinfo = DMDALocalInfo(name='info', liveness='eager') prefixes = ['y', 'x', 'f', 'b'] diff --git a/devito/petsc/solver_parameters.py b/devito/petsc/solver_parameters.py new file mode 100644 index 0000000000..7173ec9745 --- /dev/null +++ b/devito/petsc/solver_parameters.py @@ -0,0 +1,41 @@ +from petsctools import flatten_parameters +import itertools + + +# NOTE: Will be extended, the default preconditioner is not going to be 'none' +base_solve_defaults = { + 'ksp_type': 'gmres', + 'pc_type': 'none', + 'ksp_rtol': 1.e-5, + 'ksp_atol': 1.e-50, + 'ksp_divtol': 1e5, + 'ksp_max_it': int(1e4) +} + + +# Specific to linear solvers +linear_solve_defaults = { + 'snes_type': 'ksponly', + **base_solve_defaults, +} + + +def linear_solver_parameters(solver_parameters): + # Flatten parameters to support nested dictionaries + flattened = flatten_parameters(solver_parameters or {}) + processed = linear_solve_defaults.copy() + processed.update(flattened) + return processed + + +_options_prefix_counter = itertools.count() + + +def format_options_prefix(options_prefix): + # NOTE: Modified from the `OptionsManager` inside petsctools + if not options_prefix: + return f"devito_{next(_options_prefix_counter)}_" + + if options_prefix.endswith("_"): + return options_prefix + return f"{options_prefix}_" diff --git a/devito/petsc/types/object.py b/devito/petsc/types/object.py index 1484e1be81..e674a3d014 100644 --- a/devito/petsc/types/object.py +++ b/devito/petsc/types/object.py @@ -97,6 +97,14 @@ class PetscInt(PetscObject): dtype = CustomIntType('PetscInt') +class PetscScalar(PetscObject): + dtype = CustomIntType('PetscScalar') + + +class PetscBool(PetscObject): + dtype = CustomDtype('PetscBool') + + class KSP(PetscObject): """ PETSc KSP : Linear Systems Solvers. @@ -105,6 +113,14 @@ class KSP(PetscObject): dtype = CustomDtype('KSP') +class KSPType(PetscObject): + dtype = CustomDtype('KSPType') + + +class KSPNormType(PetscObject): + dtype = CustomDtype('KSPNormType') + + class CallbackSNES(PetscObject): """ PETSc SNES : Non-Linear Systems Solvers. diff --git a/devito/petsc/types/types.py b/devito/petsc/types/types.py index fa87920571..598fc658fe 100644 --- a/devito/petsc/types/types.py +++ b/devito/petsc/types/types.py @@ -1,28 +1,24 @@ import sympy - from itertools import chain from functools import cached_property from devito.tools import Reconstructable, sympy_mutex, as_tuple, frozendict from devito.tools.dtypes_lowering import dtype_mapper -from devito.petsc.utils import petsc_variables from devito.symbolics.extraction import separate_eqn, generate_targets, centre_stencil -from devito.petsc.types.equation import EssentialBC, ZeroRow, ZeroColumn from devito.types.equation import Eq from devito.operations.solve import eval_time_derivatives +from devito.petsc.utils import petsc_variables +from devito.petsc.types.equation import EssentialBC, ZeroRow, ZeroColumn + class MetaData(sympy.Function, Reconstructable): def __new__(cls, expr, **kwargs): with sympy_mutex: obj = sympy.Function.__new__(cls, expr) - obj._expr = expr + obj.expr = expr return obj - @property - def expr(self): - return self._expr - class Initialize(MetaData): pass @@ -32,105 +28,58 @@ class Finalize(MetaData): pass -class LinearSolveExpr(MetaData): +class GetArgs(MetaData): + pass + + +class SolverMetaData(MetaData): """ A symbolic expression passed through the Operator, containing the metadata - needed to execute a linear solver. Linear problems are handled with - `SNESSetType(snes, KSPONLY)`, enabling a unified interface for both - linear and nonlinear solvers. - # TODO: extend this - defaults: - - 'ksp_type': String with the name of the PETSc Krylov method. - Default is 'gmres' (Generalized Minimal Residual Method). - https://petsc.org/main/manualpages/KSP/KSPType/ - - 'pc_type': String with the name of the PETSc preconditioner. - Default is 'jacobi' (i.e diagonal scaling preconditioning). - https://petsc.org/main/manualpages/PC/PCType/ - KSP tolerances: - https://petsc.org/release/manualpages/KSP/KSPSetTolerances/ - - 'ksp_rtol': Relative convergence tolerance. Default - is 1e-5. - - 'ksp_atol': Absolute convergence for tolerance. Default - is 1e-50. - - 'ksp_divtol': Divergence tolerance, amount residual norm can - increase before `KSPConvergedDefault()` concludes - that the method is diverging. Default is 1e5. - - 'ksp_max_it': Maximum number of iterations to use. Default - is 1e4. + needed to execute the PETSc solver. """ - __rargs__ = ('expr',) __rkwargs__ = ('solver_parameters', 'field_data', 'time_mapper', - 'localinfo', 'options_prefix') - - defaults = { - 'ksp_type': 'gmres', - 'pc_type': 'jacobi', - 'ksp_rtol': 1e-5, # Relative tolerance - 'ksp_atol': 1e-50, # Absolute tolerance - 'ksp_divtol': 1e5, # Divergence tolerance - 'ksp_max_it': 1e4 # Maximum iterations - } + 'localinfo', 'user_prefix', 'formatted_prefix', + 'get_info') def __new__(cls, expr, solver_parameters=None, field_data=None, time_mapper=None, localinfo=None, - options_prefix=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) + user_prefix=None, formatted_prefix=None, + get_info=None, **kwargs): with sympy_mutex: - obj = sympy.Function.__new__(cls, expr) - - obj._expr = expr - obj._solver_parameters = solver_parameters - obj._field_data = field_data if field_data else FieldData() - obj._time_mapper = time_mapper - obj._localinfo = localinfo - obj._options_prefix = options_prefix + if isinstance(expr, tuple): + expr = sympy.Tuple(*expr) + obj = sympy.Basic.__new__(cls, expr) + + obj.expr = expr + obj.solver_parameters = frozendict(solver_parameters) + obj.field_data = field_data if field_data else FieldData() + obj.time_mapper = time_mapper + obj.localinfo = localinfo + obj.user_prefix = user_prefix + obj.formatted_prefix = formatted_prefix + obj.get_info = get_info if get_info is not None else [] return obj def __repr__(self): - return "%s(%s)" % (self.__class__.__name__, self.expr) + return f"{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) - - @property - def expr(self): - return self._expr - - @property - def field_data(self): - return self._field_data - - @property - def solver_parameters(self): - return self._solver_parameters + __hash__ = sympy.Basic.__hash__ - @property - def time_mapper(self): - return self._time_mapper - - @property - def localinfo(self): - return self._localinfo + def _hashable_content(self): + return (self.expr, self.formatted_prefix, self.solver_parameters) - @property - def options_prefix(self): - return self._options_prefix + def __eq__(self, other): + return (isinstance(other, SolverMetaData) and + self.expr == other.expr and + self.formatted_prefix == other.formatted_prefix + and self.solver_parameters == other.solver_parameters) @property def grid(self): @@ -143,9 +92,24 @@ def eval(cls, *args): func = Reconstructable._rebuild +class LinearSolverMetaData(SolverMetaData): + """ + Linear problems are handled by setting the SNESType to 'ksponly', + enabling a unified interface for both linear and nonlinear solvers. + """ + pass + + +class NonLinearSolverMetaData(SolverMetaData): + """ + TODO: Non linear solvers are not yet supported. + """ + pass + + class FieldData: """ - Metadata for a single `target` field passed to `LinearSolveExpr`. + Metadata for a single `target` field passed to `SolverMetaData`. Used to interface with PETSc SNES solvers at the IET level. Parameters @@ -170,8 +134,8 @@ def __init__(self, target=None, jacobian=None, residual=None, petsc_precision = dtype_mapper[petsc_variables['PETSC_PRECISION']] if self._target.dtype != petsc_precision: raise TypeError( - f"Your target dtype must match the precision of your " - f"PETSc configuration. " + "Your target dtype must match the precision of your " + "PETSc configuration. " f"Expected {petsc_precision}, but got {self._target.dtype}." ) self._jacobian = jacobian @@ -218,7 +182,7 @@ def targets(self): class MultipleFieldData(FieldData): """ - Metadata class passed to `LinearSolveExpr`, for mixed-field problems, + Metadata class passed to `SolverMetaData`, for mixed-field problems, where the solution vector spans multiple `targets`. Used to interface with PETSc SNES solvers at the IET level. @@ -702,7 +666,7 @@ def _build_residual(self, expr, target): class InitialGuess: """ - Metadata passed to `LinearSolveExpr` to define the initial guess + Metadata passed to `SolverExpr` to define the initial guess symbolic expressions, enforcing the initial guess to satisfy essential boundary conditions. """ diff --git a/devito/petsc/utils.py b/devito/petsc/utils.py index 3d491703e7..a0b5753255 100644 --- a/devito/petsc/utils.py +++ b/devito/petsc/utils.py @@ -1,21 +1,17 @@ import os import ctypes - from pathlib import Path -from devito.tools import memoized_func + +from devito.tools import memoized_func, filter_ordered +from devito.types import Symbol, SteppingDimension +from devito.operations.solve import eval_time_derivatives +from devito.symbolics import retrieve_functions class PetscOSError(OSError): pass -solver_mapper = { - 'gmres': 'KSPGMRES', - 'jacobi': 'PCJACOBI', - None: 'PCNONE' -} - - @memoized_func def get_petsc_dir(): petsc_dir = os.environ.get('PETSC_DIR') @@ -75,26 +71,78 @@ def get_petsc_variables(): petsc_variables = get_petsc_variables() +# TODO: Use petsctools get_petscvariables() instead? def get_petsc_type_mappings(): try: petsc_precision = petsc_variables['PETSC_PRECISION'] except KeyError: - mapper = {} + printer_mapper = {} + petsc_type_to_ctype = {} else: petsc_scalar = 'PetscScalar' # TODO: Check to see whether Petsc is compiled with # 32-bit or 64-bit integers - mapper = {ctypes.c_int: 'PetscInt'} + printer_mapper = {ctypes.c_int: 'PetscInt'} if petsc_precision == 'single': - mapper[ctypes.c_float] = petsc_scalar + printer_mapper[ctypes.c_float] = petsc_scalar elif petsc_precision == 'double': - mapper[ctypes.c_double] = petsc_scalar - return mapper + printer_mapper[ctypes.c_double] = petsc_scalar + + # Used to construct ctypes.Structures that wrap PETSc objects + petsc_type_to_ctype = {v: k for k, v in printer_mapper.items()} + # Add other PETSc types + petsc_type_to_ctype.update({ + 'KSPType': ctypes.c_char_p, + 'KSPConvergedReason': petsc_type_to_ctype['PetscInt'], + 'KSPNormType': petsc_type_to_ctype['PetscInt'], + }) + return printer_mapper, petsc_type_to_ctype + +petsc_type_mappings, petsc_type_to_ctype = get_petsc_type_mappings() -petsc_type_mappings = get_petsc_type_mappings() petsc_languages = ['petsc'] + + +def get_funcs(exprs): + funcs = [ + f for e in exprs + for f in retrieve_functions(eval_time_derivatives(e.lhs - e.rhs)) + ] + return filter_ordered(funcs) + + +def generate_time_mapper(funcs): + """ + Replace time indices with `Symbols` in expressions used within + PETSc callback functions. These symbols are Uxreplaced at the IET + level to align with the `TimeDimension` and `ModuloDimension` objects + present in the initial lowering. + NOTE: All functions used in PETSc callback functions are attached to + the `SolverMetaData` object, which is passed through the initial lowering + (and subsequently dropped and replaced with calls to run the solver). + Therefore, the appropriate time loop will always be correctly generated inside + the main kernel. + Examples + -------- + >>> funcs = [ + >>> f1(t + dt, x, y), + >>> g1(t + dt, x, y), + >>> g2(t, x, y), + >>> f1(t, x, y) + >>> ] + >>> generate_time_mapper(funcs) + {t + dt: tau0, t: tau1} + """ + time_indices = list({ + i if isinstance(d, SteppingDimension) else d + 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 dict(zip(time_indices, tau_symbs)) diff --git a/devito/symbolics/extended_sympy.py b/devito/symbolics/extended_sympy.py index 4016a88466..19a6640f99 100644 --- a/devito/symbolics/extended_sympy.py +++ b/devito/symbolics/extended_sympy.py @@ -916,3 +916,6 @@ def rfunc(func, item, *args): min: Min, max: Max, } + + +Null = Macro('NULL') diff --git a/requirements.txt b/requirements.txt index 4c90867e92..e4781328d3 100644 --- a/requirements.txt +++ b/requirements.txt @@ -8,4 +8,5 @@ codepy>=2019.1,<2025 multidict<6.3 anytree>=2.4.3,<=2.13.0 cloudpickle<3.1.2 -packaging<25.1 \ No newline at end of file +packaging<25.1 +petsctools<=2025.1 diff --git a/tests/test_petsc.py b/tests/test_petsc.py index 207294dbaa..28f7c04025 100644 --- a/tests/test_petsc.py +++ b/tests/test_petsc.py @@ -2,7 +2,7 @@ import numpy as np import os -from collections import OrderedDict +import re from conftest import skipif from devito import (Grid, Function, TimeFunction, Eq, Operator, @@ -14,21 +14,54 @@ from devito.passes.iet.languages.C import CDataManager from devito.petsc.types import (DM, Mat, Vec, PetscMPIInt, KSP, PC, KSPConvergedReason, PETScArray, - LinearSolveExpr, FieldData, MultipleFieldData, + FieldData, MultipleFieldData, SubMatrixBlock) from devito.petsc.solve import PETScSolve, EssentialBC from devito.petsc.iet.nodes import Expression from devito.petsc.initialize import PetscInitialize from devito.petsc.logging import PetscSummary +from devito.petsc.solver_parameters import linear_solve_defaults + + +@pytest.fixture(scope='session') +def command_line(): + + # Random prefixes to validate command line argument parsing + prefix = ( + 'd17weqroeg', 'riabfodkj5', 'fir8o3lsak', + 'zwejklqn25', 'qtr2vfvwiu') + + petsc_option = ( + ('ksp_rtol',), + ('ksp_rtol', 'ksp_atol'), + ('ksp_rtol', 'ksp_atol', 'ksp_divtol', 'ksp_max_it'), + ('ksp_type',), + ('ksp_divtol', 'ksp_type') + ) + value = ( + (1e-8,), + (1e-11, 1e-15), + (1e-3, 1e-10, 50000, 2000), + ('cg',), + (22000, 'richardson'), + ) + argv = [] + expected = {} + for p, opt, val in zip(prefix, petsc_option, value, strict=True): + for o, v in zip(opt, val, strict=True): + argv.extend([f'-{p}_{o}', str(v)]) + expected[p] = zip(opt, val) + return argv, expected @pytest.fixture(scope='session', autouse=True) -def petsc_initialization(): +def petsc_initialization(command_line): + argv, _ = command_line # TODO: Temporary workaround until PETSc is automatically # initialized configuration['compiler'] = 'custom' os.environ['CC'] = 'mpicc' - PetscInitialize() + PetscInitialize(argv) @skipif('petsc') @@ -158,16 +191,18 @@ def test_multiple_petsc_solves(): eqn1 = Eq(f1.laplace, g1) eqn2 = Eq(f2.laplace, g2) - petsc1 = PETScSolve(eqn1, f1) - petsc2 = PETScSolve(eqn2, f2) + petsc1 = PETScSolve(eqn1, f1, options_prefix='pde1') + petsc2 = PETScSolve(eqn2, f2, options_prefix='pde2') with switchconfig(language='petsc'): op = Operator([petsc1, petsc2], opt='noop') callable_roots = [meta_call.root for meta_call in op._func_table.values()] - # One FormRHS, MatShellMult, FormFunction, PopulateMatContext per solve - assert len(callable_roots) == 8 + # One FormRHS, MatShellMult, FormFunction, PopulateMatContext, SetPetscOptions + # and ClearPetscOptions per solve. + # TODO: Some efuncs are not reused where reuse is possible — investigate. + assert len(callable_roots) == 12 @skipif('petsc') @@ -204,24 +239,6 @@ def test_petsc_cast(): '(PetscScalar (*)[info.gym][info.gxm]) x_f3_vec;' in str(op3.ccode) -@skipif('petsc') -def test_LinearSolveExpr(): - - grid = Grid((2, 2), dtype=np.float64) - - 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, field_data=FieldData(target=f)) - - # Check the solver parameters - assert linsolveexpr.solver_parameters == \ - {'ksp_type': 'gmres', 'pc_type': 'jacobi', 'ksp_rtol': 1e-05, - 'ksp_atol': 1e-50, 'ksp_divtol': 100000.0, 'ksp_max_it': 10000} - - @skipif('petsc') def test_dmda_create(): @@ -786,8 +803,10 @@ def test_coupled_vs_non_coupled(self, eq1, eq2, so): # Solving for multiple fields within the same matrix system requires # less callback functions than solving them separately. - assert len(callbacks1) == 8 - assert len(callbacks2) == 6 + # TODO: As noted in the other test, some efuncs are not reused + # where reuse is possible, investigate. + assert len(callbacks1) == 12 + assert len(callbacks2) == 8 # Check field_data type field0 = petsc1.rhs.field_data @@ -1313,7 +1332,7 @@ def define(self, dimensions): '+ r1*a0[ix + 2][iy + 3]))*o0->h_x*o0->h_y;' in str(J00) # J00 and J11 are semantically identical so check efunc reuse - assert len(op._func_table.values()) == 7 + assert len(op._func_table.values()) == 9 # J00_MatMult0 is reused (in replace of J11_MatMult0) create = op._func_table['MatCreateSubMatrices0'].root assert 'MatShellSetOperation(submat_arr[0],' \ @@ -1416,27 +1435,32 @@ def test_logging(self, log_level): entry0 = petsc_summary.get_entry('section0', 'poisson') entry1 = petsc_summary[('section0', 'poisson')] assert entry0 == entry1 - assert str(entry0) == \ - 'PetscEntry(KSPGetIterationNumber=16, SNESGetIterationNumber=1)' - assert entry0.SNESGetIterationNumber == 1 - snesiters0 = petsc_summary.SNESGetIterationNumber - snesiters1 = petsc_summary['SNESGetIterationNumber'] + snesits0 = petsc_summary.SNESGetIterationNumber + snesits1 = petsc_summary['SNESGetIterationNumber'] # Check case insensitive key access - snesiters2 = petsc_summary['snesgetiterationnumber'] - snesiters3 = petsc_summary['SNESgetiterationNumber'] + snesits2 = petsc_summary['snesgetiterationnumber'] + snesits3 = petsc_summary['SNESgetiterationNumber'] - assert snesiters0 == snesiters1 == snesiters2 == snesiters3 + assert snesits0 == snesits1 == snesits2 == snesits3 - assert isinstance(snesiters0, OrderedDict) - assert len(snesiters0) == 1 - key, value = next(iter(snesiters0.items())) + assert len(snesits0) == 1 + key, value = next(iter(snesits0.items())) assert str(key) == "PetscKey(name='section0', options_prefix='poisson')" assert value == 1 + # Test logging KSPGetTolerances. Since no overrides have been applied, + # the tolerances should match the default linear values. + tols = entry0.KSPGetTolerances + assert tols['rtol'] == linear_solve_defaults['ksp_rtol'] + assert tols['atol'] == linear_solve_defaults['ksp_atol'] + assert tols['divtol'] == linear_solve_defaults['ksp_divtol'] + assert tols['max_it'] == linear_solve_defaults['ksp_max_it'] + @skipif('petsc') - def test_logging_multiple_solves(self): + @pytest.mark.parametrize('log_level', ['PERF', 'DEBUG']) + def test_logging_multiple_solves(self, log_level): grid = Grid(shape=(11, 11), dtype=np.float64) functions = [Function(name=n, grid=grid, space_order=2) @@ -1452,7 +1476,7 @@ def test_logging_multiple_solves(self): solver1 = PETScSolve(eq1, target=g, options_prefix='poisson1') solver2 = PETScSolve(eq2, target=h, options_prefix='poisson2') - with switchconfig(language='petsc', log_level='DEBUG'): + with switchconfig(language='petsc', log_level=log_level): op = Operator([solver1, solver2]) summary = op.apply() @@ -1463,11 +1487,6 @@ def test_logging_multiple_solves(self): entry1 = petsc_summary.get_entry('section0', 'poisson1') entry2 = petsc_summary.get_entry('section1', 'poisson2') - assert str(entry1) == \ - 'PetscEntry(KSPGetIterationNumber=16, SNESGetIterationNumber=1)' - assert str(entry2) == \ - 'PetscEntry(KSPGetIterationNumber=1, SNESGetIterationNumber=1)' - assert len(petsc_summary.KSPGetIterationNumber) == 2 assert len(petsc_summary.SNESGetIterationNumber) == 2 @@ -1481,3 +1500,579 @@ def test_logging_multiple_solves(self): assert entry1['SNESGetIterationNumber'] == 1 # Case insensitive key access assert entry1['kspgetiterationnumber'] == 16 + + @skipif('petsc') + @pytest.mark.parametrize('log_level', ['PERF', 'DEBUG']) + def test_logging_user_prefixes(self, log_level): + """ + Verify that `PetscSummary` uses the user provided `options_prefix` when given. + """ + grid = Grid(shape=(11, 11), dtype=np.float64) + + functions = [Function(name=n, grid=grid, space_order=2) + for n in ['e', 'f', 'g', 'h']] + e, f, g, h = functions + + pde1 = Eq(e.laplace, f) + pde2 = Eq(g.laplace, h) + + petsc1 = PETScSolve(pde1, target=e, options_prefix='pde1') + petsc2 = PETScSolve(pde2, target=g, options_prefix='pde2') + + with switchconfig(language='petsc', log_level=log_level): + op = Operator([petsc1, petsc2]) + summary = op.apply() + + petsc_summary = summary.petsc + + # Check that the prefix is correctly set in the PetscSummary + key_strings = [f"{key.name}:{key.options_prefix}" for key in petsc_summary.keys()] + assert set(key_strings) == {"section0:pde1", "section1:pde2"} + + @skipif('petsc') + @pytest.mark.parametrize('log_level', ['PERF', 'DEBUG']) + def test_logging_default_prefixes(self, log_level): + """ + Verify that `PetscSummary` uses the default options prefix + provided by Devito if no user `options_prefix` is specified. + """ + grid = Grid(shape=(11, 11), dtype=np.float64) + + functions = [Function(name=n, grid=grid, space_order=2) + for n in ['e', 'f', 'g', 'h']] + e, f, g, h = functions + + pde1 = Eq(e.laplace, f) + pde2 = Eq(g.laplace, h) + + petsc1 = PETScSolve(pde1, target=e) + petsc2 = PETScSolve(pde2, target=g) + + with switchconfig(language='petsc', log_level=log_level): + op = Operator([petsc1, petsc2]) + summary = op.apply() + + petsc_summary = summary.petsc + + # Users should set a custom options_prefix if they want logging; otherwise, + # the default automatically generated prefix is used in the `PetscSummary`. + assert all(re.fullmatch(r"devito_\d+_", k.options_prefix) for k in petsc_summary) + + +class TestSolverParameters: + + @skipif('petsc') + def setup_class(self): + """ + Setup grid, functions and equations shared across + tests in this class + """ + grid = Grid(shape=(11, 11), dtype=np.float64) + self.e, self.f, self.g, self.h = [ + Function(name=n, grid=grid, space_order=2) + for n in ['e', 'f', 'g', 'h'] + ] + self.eq1 = Eq(self.e.laplace, self.f) + self.eq2 = Eq(self.g.laplace, self.h) + + @skipif('petsc') + def test_different_solver_params(self): + # Explicitly set the solver parameters + solver1 = PETScSolve( + self.eq1, target=self.e, solver_parameters={'ksp_rtol': '1e-10'} + ) + # Use solver parameter defaults + solver2 = PETScSolve(self.eq2, target=self.g) + + with switchconfig(language='petsc'): + op = Operator([solver1, solver2]) + + assert 'SetPetscOptions0' in op._func_table + assert 'SetPetscOptions1' in op._func_table + + assert '_ksp_rtol","1e-10"' \ + in str(op._func_table['SetPetscOptions0'].root) + + assert '_ksp_rtol","1e-05"' \ + in str(op._func_table['SetPetscOptions1'].root) + + @skipif('petsc') + def test_options_prefix(self): + solver1 = PETScSolve(self.eq1, self.e, + solver_parameters={'ksp_rtol': '1e-10'}, + options_prefix='poisson1') + solver2 = PETScSolve(self.eq2, self.g, + solver_parameters={'ksp_rtol': '1e-12'}, + options_prefix='poisson2') + + with switchconfig(language='petsc'): + op = Operator([solver1, solver2]) + + # Check the options prefix has been correctly set for each snes solver + assert 'PetscCall(SNESSetOptionsPrefix(snes0,"poisson1_"));' in str(op) + assert 'PetscCall(SNESSetOptionsPrefix(snes1,"poisson2_"));' in str(op) + + # Test the options prefix has be correctly applied to the solver options + assert 'PetscCall(PetscOptionsSetValue(NULL,"-poisson1_ksp_rtol","1e-10"));' \ + in str(op._func_table['SetPetscOptions0'].root) + + assert 'PetscCall(PetscOptionsSetValue(NULL,"-poisson2_ksp_rtol","1e-12"));' \ + in str(op._func_table['SetPetscOptions1'].root) + + @skipif('petsc') + def test_options_no_value(self): + """ + Test solver parameters that do not require a value, such as + `snes_view` and `ksp_view`. + """ + solver = PETScSolve( + self.eq1, target=self.e, solver_parameters={'snes_view': None}, + options_prefix='solver1' + ) + with switchconfig(language='petsc'): + op = Operator(solver) + op.apply() + + assert 'PetscCall(PetscOptionsSetValue(NULL,"-solver1_snes_view",NULL));' \ + in str(op._func_table['SetPetscOptions0'].root) + + @skipif('petsc') + @pytest.mark.parametrize('log_level', ['PERF', 'DEBUG']) + def test_tolerances(self, log_level): + params = { + 'ksp_rtol': 1e-12, + 'ksp_atol': 1e-20, + 'ksp_divtol': 1e3, + 'ksp_max_it': 100 + } + solver = PETScSolve( + self.eq1, target=self.e, solver_parameters=params, + options_prefix='solver' + ) + + with switchconfig(language='petsc', log_level=log_level): + op = Operator(solver) + tmp = op.apply() + + petsc_summary = tmp.petsc + entry = petsc_summary.get_entry('section0', 'solver') + tolerances = entry.KSPGetTolerances + + # Test that the tolerances have been set correctly and therefore + # appear as expected in the `PetscSummary`. + assert tolerances['rtol'] == params['ksp_rtol'] + assert tolerances['atol'] == params['ksp_atol'] + assert tolerances['divtol'] == params['ksp_divtol'] + assert tolerances['max_it'] == params['ksp_max_it'] + + @skipif('petsc') + def test_clearing_options(self): + # Explicitly set the solver parameters + solver1 = PETScSolve( + self.eq1, target=self.e, solver_parameters={'ksp_rtol': '1e-10'} + ) + # Use the solver parameter defaults + solver2 = PETScSolve(self.eq2, target=self.g) + + with switchconfig(language='petsc'): + op = Operator([solver1, solver2]) + + assert 'ClearPetscOptions0' in op._func_table + assert 'ClearPetscOptions1' in op._func_table + + @skipif('petsc') + def test_error_if_same_prefix(self): + """ + Test an error is raised if the same options prefix is used + for two different solvers within the same Operator. + """ + solver1 = PETScSolve( + self.eq1, target=self.e, options_prefix='poisson', + solver_parameters={'ksp_rtol': '1e-10'} + ) + solver2 = PETScSolve( + self.eq2, target=self.g, options_prefix='poisson', + solver_parameters={'ksp_rtol': '1e-12'} + ) + with switchconfig(language='petsc'): + with pytest.raises(ValueError): + Operator([solver1, solver2]) + + @skipif('petsc') + @pytest.mark.parametrize('log_level', ['PERF', 'DEBUG']) + def test_multiple_operators(self, log_level): + """ + Verify that solver parameters are set correctly when multiple Operators + are created with PETScSolve instances sharing the same options_prefix. + + Note: Using the same options_prefix within a single Operator is not allowed + (see previous test), but the same prefix can be used across + different Operators (although not advised). + """ + # Create two PETScSolve instances with the same options_prefix + solver1 = PETScSolve( + self.eq1, target=self.e, options_prefix='poisson', + solver_parameters={'ksp_rtol': '1e-10'} + ) + solver2 = PETScSolve( + self.eq2, target=self.g, options_prefix='poisson', + solver_parameters={'ksp_rtol': '1e-12'} + ) + with switchconfig(language='petsc', log_level=log_level): + op1 = Operator(solver1) + op2 = Operator(solver2) + summary1 = op1.apply() + summary2 = op2.apply() + + petsc_summary1 = summary1.petsc + entry1 = petsc_summary1.get_entry('section0', 'poisson') + + petsc_summary2 = summary2.petsc + entry2 = petsc_summary2.get_entry('section0', 'poisson') + + assert entry1.KSPGetTolerances['rtol'] == 1e-10 + assert entry2.KSPGetTolerances['rtol'] == 1e-12 + + @skipif('petsc') + @pytest.mark.parametrize('log_level', ['PERF', 'DEBUG']) + def test_command_line_priority_tols_1(self, command_line, log_level): + """ + Test solver tolerances specifed via the command line + take precedence over those specified in the defaults. + """ + prefix = 'd17weqroeg' + _, expected = command_line + + solver1 = PETScSolve( + self.eq1, target=self.e, + options_prefix=prefix + ) + with switchconfig(language='petsc', log_level=log_level): + op = Operator(solver1) + summary = op.apply() + + petsc_summary = summary.petsc + entry = petsc_summary.get_entry('section0', prefix) + for opt, val in expected[prefix]: + assert entry.KSPGetTolerances[opt.removeprefix('ksp_')] == val + + @skipif('petsc') + @pytest.mark.parametrize('log_level', ['PERF', 'DEBUG']) + def test_command_line_priority_tols_2(self, command_line, log_level): + prefix = 'riabfodkj5' + _, expected = command_line + + solver1 = PETScSolve( + self.eq1, target=self.e, + options_prefix=prefix + ) + with switchconfig(language='petsc', log_level=log_level): + op = Operator(solver1) + summary = op.apply() + + petsc_summary = summary.petsc + entry = petsc_summary.get_entry('section0', prefix) + for opt, val in expected[prefix]: + assert entry.KSPGetTolerances[opt.removeprefix('ksp_')] == val + + @skipif('petsc') + @pytest.mark.parametrize('log_level', ['PERF', 'DEBUG']) + def test_command_line_priority_tols3(self, command_line, log_level): + """ + Test solver tolerances specifed via the command line + take precedence over those specified by the `solver_parameters` dict. + """ + prefix = 'fir8o3lsak' + _, expected = command_line + + # Set solver parameters that differ both from the defaults and from the + # values provided on the command line for this prefix (see the `command_line` + # fixture). + params = { + 'ksp_rtol': 1e-13, + 'ksp_atol': 1e-35, + 'ksp_divtol': 300000, + 'ksp_max_it': 500 + } + + solver1 = PETScSolve( + self.eq1, target=self.e, + solver_parameters=params, + options_prefix=prefix + ) + with switchconfig(language='petsc', log_level=log_level): + op = Operator(solver1) + summary = op.apply() + + petsc_summary = summary.petsc + entry = petsc_summary.get_entry('section0', prefix) + for opt, val in expected[prefix]: + assert entry.KSPGetTolerances[opt.removeprefix('ksp_')] == val + + @skipif('petsc') + @pytest.mark.parametrize('log_level', ['PERF', 'DEBUG']) + def test_command_line_priority_ksp_type(self, command_line, log_level): + """ + Test the solver parameter 'ksp_type' specified via the command line + take precedence over the one specified in the `solver_parameters` dict. + """ + prefix = 'zwejklqn25' + _, expected = command_line + + # Set `ksp_type`` in the solver parameters, which should be overridden + # by the command line value (which is set to `cg` - + # see the `command_line` fixture). + params = {'ksp_type': 'richardson'} + + solver1 = PETScSolve( + self.eq1, target=self.e, + solver_parameters=params, + options_prefix=prefix + ) + with switchconfig(language='petsc', log_level=log_level): + op = Operator(solver1) + summary = op.apply() + + petsc_summary = summary.petsc + entry = petsc_summary.get_entry('section0', prefix) + for _, val in expected[prefix]: + assert entry.KSPGetType == val + assert not entry.KSPGetType == params['ksp_type'] + + @skipif('petsc') + def test_command_line_priority_ccode(self, command_line): + """ + Verify that if an option is set via the command line, + the corresponding entry in `linear_solve_defaults` or `solver_parameters` + is not set or cleared in the generated code. (The command line option + will have already been set in the global PetscOptions database + during PetscInitialize().) + """ + prefix = 'qtr2vfvwiu' + + solver = PETScSolve( + self.eq1, target=self.e, + # Specify a solver parameter that is not set via the + # command line (see the `command_line` fixture for this prefix). + solver_parameters={'ksp_rtol': '1e-10'}, + options_prefix=prefix + ) + with switchconfig(language='petsc'): + op = Operator(solver) + + set_options_callback = str(op._func_table['SetPetscOptions0'].root) + clear_options_callback = str(op._func_table['ClearPetscOptions0'].root) + + # Check that the `ksp_rtol` option IS set and cleared explicitly + # since it is NOT set via the command line. + assert f'PetscOptionsSetValue(NULL,"-{prefix}_ksp_rtol","1e-10")' \ + in set_options_callback + assert f'PetscOptionsClearValue(NULL,"-{prefix}_ksp_rtol")' \ + in clear_options_callback + + # Check that the `ksp_divtol` and `ksp_type` options are NOT set + # or cleared explicitly since they ARE set via the command line. + assert f'PetscOptionsSetValue(NULL,"-{prefix}_div_tol",' \ + not in set_options_callback + assert f'PetscOptionsSetValue(NULL,"-{prefix}_ksp_type",' \ + not in set_options_callback + assert f'PetscOptionsClearValue(NULL,"-{prefix}_div_tol"));' \ + not in clear_options_callback + assert f'PetscOptionsClearValue(NULL,"-{prefix}_ksp_type"));' \ + not in clear_options_callback + + # Check that options specifed by the `linear_solver_defaults` + # are still set and cleared + assert f'PetscOptionsSetValue(NULL,"-{prefix}_ksp_atol",' \ + in set_options_callback + assert f'PetscOptionsClearValue(NULL,"-{prefix}_ksp_atol"));' \ + in clear_options_callback + + +class TestHashing: + + @skipif('petsc') + def test_solveexpr(self): + grid = Grid(shape=(11, 11), dtype=np.float64) + functions = [Function(name=n, grid=grid, space_order=2) + for n in ['e', 'f']] + e, f = functions + eq = Eq(e.laplace, f) + + # Two PETScSolve instances with different options_prefix values + # should hash differently. + petsc1 = PETScSolve(eq, target=e, options_prefix='poisson1') + petsc2 = PETScSolve(eq, target=e, options_prefix='poisson2') + + assert hash(petsc1.rhs) != hash(petsc2.rhs) + assert petsc1.rhs != petsc2.rhs + + # Two PETScSolve instances with the same options_prefix but + # different solver parameters should hash differently. + petsc3 = PETScSolve( + eq, target=e, solver_parameters={'ksp_type': 'cg'}, + options_prefix='poisson3' + ) + petsc4 = PETScSolve( + eq, target=e, solver_parameters={'ksp_type': 'richardson'}, + options_prefix='poisson3' + ) + assert hash(petsc3.rhs) != hash(petsc4.rhs) + + +class TestGetInfo: + """ + Test the `get_info` optional argument to `PETScSolve`. + + This argument can be used independently of the `log_level` to retrieve + specific information about the solve, such as the number of KSP + iterations to converge. + """ + @skipif('petsc') + def setup_class(self): + """ + Setup grid, functions and equations shared across + tests in this class + """ + grid = Grid(shape=(11, 11), dtype=np.float64) + self.e, self.f, self.g, self.h = [ + Function(name=n, grid=grid, space_order=2) + for n in ['e', 'f', 'g', 'h'] + ] + self.eq1 = Eq(self.e.laplace, self.f) + self.eq2 = Eq(self.g.laplace, self.h) + + @skipif('petsc') + def test_get_info(self): + get_info = ['kspgetiterationnumber', 'snesgetiterationnumber'] + petsc = PETScSolve( + self.eq1, target=self.e, options_prefix='pde1', get_info=get_info + ) + with switchconfig(language='petsc'): + op = Operator(petsc) + summary = op.apply() + + petsc_summary = summary.petsc + entry = petsc_summary.get_entry('section0', 'pde1') + + # Verify that the entry contains only the requested info + # (since logging is not set) + assert len(entry) == 2 + assert hasattr(entry, "KSPGetIterationNumber") + assert hasattr(entry, "SNESGetIterationNumber") + + @skipif('petsc') + @pytest.mark.parametrize('log_level', ['PERF', 'DEBUG']) + def test_get_info_with_logging(self, log_level): + """ + Test that `get_info` works correctly when logging is enabled. + """ + get_info = ['kspgetiterationnumber'] + petsc = PETScSolve( + self.eq1, target=self.e, options_prefix='pde1', get_info=get_info + ) + with switchconfig(language='petsc', log_level=log_level): + op = Operator(petsc) + summary = op.apply() + + petsc_summary = summary.petsc + entry = petsc_summary.get_entry('section0', 'pde1') + + # With logging enabled, the entry should include both the + # requested KSP iteration number and additional PETSc info + # (e.g., SNES iteration count logged at PERF/DEBUG). + assert len(entry) > 1 + assert hasattr(entry, "KSPGetIterationNumber") + assert hasattr(entry, "SNESGetIterationNumber") + + @skipif('petsc') + def test_different_solvers(self): + """ + Test that `get_info` works correctly when multiple solvers are used + within the same Operator. + """ + # Create two PETScSolve instances with different get_info arguments + + get_info_1 = ['kspgetiterationnumber'] + get_info_2 = ['snesgetiterationnumber'] + + solver1 = PETScSolve( + self.eq1, target=self.e, options_prefix='pde1', get_info=get_info_1 + ) + solver2 = PETScSolve( + self.eq2, target=self.g, options_prefix='pde2', get_info=get_info_2 + ) + with switchconfig(language='petsc'): + op = Operator([solver1, solver2]) + summary = op.apply() + + petsc_summary = summary.petsc + + assert len(petsc_summary) == 2 + assert len(petsc_summary.KSPGetIterationNumber) == 1 + assert len(petsc_summary.SNESGetIterationNumber) == 1 + + entry1 = petsc_summary.get_entry('section0', 'pde1') + entry2 = petsc_summary.get_entry('section1', 'pde2') + + assert hasattr(entry1, "KSPGetIterationNumber") + assert not hasattr(entry1, "SNESGetIterationNumber") + + assert not hasattr(entry2, "KSPGetIterationNumber") + assert hasattr(entry2, "SNESGetIterationNumber") + + @skipif('petsc') + def test_case_insensitive(self): + """ + Test that `get_info` is case insensitive + """ + # Create a list with mixed cases + get_info = ['KSPGetIterationNumber', 'snesgetiterationnumber'] + petsc = PETScSolve( + self.eq1, target=self.e, options_prefix='pde1', get_info=get_info + ) + with switchconfig(language='petsc'): + op = Operator(petsc) + summary = op.apply() + + petsc_summary = summary.petsc + entry = petsc_summary.get_entry('section0', 'pde1') + + assert hasattr(entry, "KSPGetIterationNumber") + assert hasattr(entry, "SNESGetIterationNumber") + + @skipif('petsc') + def test_get_ksp_type(self): + """ + Test that `get_info` can retrieve the KSP type as + a string. + """ + get_info = ['kspgettype'] + solver1 = PETScSolve( + self.eq1, target=self.e, options_prefix='poisson1', get_info=get_info + ) + solver2 = PETScSolve( + self.eq1, target=self.e, options_prefix='poisson2', + solver_parameters={'ksp_type': 'cg'}, get_info=get_info + ) + with switchconfig(language='petsc'): + op = Operator([solver1, solver2]) + summary = op.apply() + + petsc_summary = summary.petsc + entry1 = petsc_summary.get_entry('section0', 'poisson1') + entry2 = petsc_summary.get_entry('section1', 'poisson2') + + assert hasattr(entry1, "KSPGetType") + # Check the type matches the default in linear_solve_defaults + # since it has not been overridden + assert entry1.KSPGetType == linear_solve_defaults['ksp_type'] + assert entry1['KSPGetType'] == linear_solve_defaults['ksp_type'] + assert entry1['kspgettype'] == linear_solve_defaults['ksp_type'] + + # Test that the KSP type default is correctly overridden by the + # solver_parameters dictionary passed to solver2 + assert hasattr(entry2, "KSPGetType") + assert entry2.KSPGetType == 'cg' + assert entry2['KSPGetType'] == 'cg' + assert entry2['kspgettype'] == 'cg' From a19f13614600cf7c28be865df483f8e132b99605 Mon Sep 17 00:00:00 2001 From: ZoeLeibowitz Date: Fri, 19 Sep 2025 10:17:22 +0100 Subject: [PATCH 81/84] workflows: Fix segfault in parallel tests --- .github/workflows/pytest-petsc.yml | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/.github/workflows/pytest-petsc.yml b/.github/workflows/pytest-petsc.yml index b3ca034500..5851c69594 100644 --- a/.github/workflows/pytest-petsc.yml +++ b/.github/workflows/pytest-petsc.yml @@ -11,12 +11,10 @@ on: branches: - main - petsc - - biharmonic pull_request: branches: - main - petsc - - biharmonic jobs: pytest: @@ -78,7 +76,7 @@ jobs: - name: Test with pytest - parallel run: | - ${{ env.RUN_CMD }} python3 -m pytest --cov --cov-config=.coveragerc --cov-report=xml -m parallel ${{ env.TESTS }} + ${{ env.RUN_CMD }} mpiexec -n 1 pytest -m parallel --cov --cov-config=.coveragerc --cov-report=xml ${{ env.TESTS }} - name: Test examples run: | From cdb83ef6960d3e0723c30c25cbe2f59c42173682 Mon Sep 17 00:00:00 2001 From: ZoeLeibowitz <90093761+ZoeLeibowitz@users.noreply.github.com> Date: Mon, 22 Sep 2025 12:23:49 +0100 Subject: [PATCH 82/84] PETSc clean up (#2740) * misc/compiler: Rename/separate files, clean up, add lower_petsc_symbols and change PETScSolve to petscsolve --- conftest.py | 2 +- devito/core/cpu.py | 8 + devito/ir/iet/algorithms.py | 2 +- devito/ir/iet/visitors.py | 2 +- devito/passes/iet/languages/C.py | 6 +- devito/petsc/config.py | 86 + devito/petsc/iet/builder.py | 341 +++ devito/petsc/iet/callbacks.py | 1132 ++++++++++ devito/petsc/iet/logging.py | 2 +- devito/petsc/iet/nodes.py | 9 + devito/petsc/iet/passes.py | 269 +-- devito/petsc/iet/routines.py | 1877 ----------------- devito/petsc/iet/solve.py | 156 ++ devito/petsc/iet/time_dependence.py | 197 ++ devito/petsc/iet/type_builder.py | 252 +++ devito/petsc/iet/utils.py | 73 - devito/petsc/initialize.py | 1 + devito/petsc/logging.py | 11 +- devito/petsc/solve.py | 82 +- devito/petsc/solver_parameters.py | 3 +- devito/petsc/types/__init__.py | 3 +- devito/petsc/types/array.py | 2 +- devito/petsc/types/equation.py | 6 +- devito/petsc/types/macros.py | 1 - devito/petsc/types/{types.py => metadata.py} | 23 +- devito/petsc/types/modes.py | 16 + devito/petsc/types/object.py | 36 +- devito/petsc/utils.py | 148 -- examples/petsc/Poisson/01_poisson.py | 4 +- examples/petsc/Poisson/02_laplace.py | 4 +- examples/petsc/Poisson/03_poisson.py | 4 +- examples/petsc/Poisson/04_poisson.py | 4 +- examples/petsc/cfd/01_navierstokes.py | 8 +- examples/petsc/petsc_test.py | 4 +- examples/petsc/random/01_helmholtz.py | 4 +- examples/petsc/random/02_biharmonic.py | 4 +- .../petsc/seismic/01_staggered_acoustic.py | 14 +- tests/test_petsc.py | 647 +++--- 38 files changed, 2882 insertions(+), 2561 deletions(-) create mode 100644 devito/petsc/config.py create mode 100644 devito/petsc/iet/builder.py create mode 100644 devito/petsc/iet/callbacks.py delete mode 100644 devito/petsc/iet/routines.py create mode 100644 devito/petsc/iet/solve.py create mode 100644 devito/petsc/iet/time_dependence.py create mode 100644 devito/petsc/iet/type_builder.py delete mode 100644 devito/petsc/iet/utils.py rename devito/petsc/types/{types.py => metadata.py} (96%) create mode 100644 devito/petsc/types/modes.py delete mode 100644 devito/petsc/utils.py diff --git a/conftest.py b/conftest.py index c0339f722b..0ed5b9ad78 100644 --- a/conftest.py +++ b/conftest.py @@ -14,7 +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 PetscOSError, get_petsc_dir +from devito.petsc.config import PetscOSError, get_petsc_dir try: from mpi4py import MPI # noqa diff --git a/devito/core/cpu.py b/devito/core/cpu.py index a93d3a3fef..eaa9266dc8 100644 --- a/devito/core/cpu.py +++ b/devito/core/cpu.py @@ -14,6 +14,8 @@ check_stability, PetscTarget) from devito.tools import timed_pass +from devito.petsc.iet.passes import lower_petsc_symbols + __all__ = ['Cpu64NoopCOperator', 'Cpu64NoopOmpOperator', 'Cpu64AdvCOperator', 'Cpu64AdvOmpOperator', 'Cpu64FsgCOperator', 'Cpu64FsgOmpOperator', 'Cpu64CustomOperator', 'Cpu64CustomCXXOperator', 'Cpu64AdvCXXOperator', @@ -143,6 +145,9 @@ def _specialize_iet(cls, graph, **kwargs): # Symbol definitions cls._Target.DataManager(**kwargs).process(graph) + # Lower PETSc symbols + lower_petsc_symbols(graph, **kwargs) + return graph @@ -222,6 +227,9 @@ def _specialize_iet(cls, graph, **kwargs): # Symbol definitions cls._Target.DataManager(**kwargs).process(graph) + # Lower PETSc symbols + lower_petsc_symbols(graph, **kwargs) + # Linearize n-dimensional Indexeds linearize(graph, **kwargs) diff --git a/devito/ir/iet/algorithms.py b/devito/ir/iet/algorithms.py index 52f48e28b1..01e3b4976e 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 MetaData -from devito.petsc.iet.utils import petsc_iet_mapper +from devito.petsc.iet.nodes import petsc_iet_mapper __all__ = ['iet_build'] diff --git a/devito/ir/iet/visitors.py b/devito/ir/iet/visitors.py index 1179cd8936..a72772b854 100644 --- a/devito/ir/iet/visitors.py +++ b/devito/ir/iet/visitors.py @@ -1068,7 +1068,7 @@ class FindSymbols(LazyVisitor[Any, list[Any], None]): Drive the search. Accepted: - `symbolics`: Collect all AbstractFunction objects, default - `basics`: Collect all Basic objects - - `abstractsymbols`: Collect all AbstractSymbol objects + - `symbols`: Collect all AbstractSymbol objects - `dimensions`: Collect all Dimensions - `indexeds`: Collect all Indexed objects - `indexedbases`: Collect all IndexedBase objects diff --git a/devito/passes/iet/languages/C.py b/devito/passes/iet/languages/C.py index 7822bf1680..6f2850c3a9 100644 --- a/devito/passes/iet/languages/C.py +++ b/devito/passes/iet/languages/C.py @@ -7,7 +7,8 @@ from devito.passes.iet.langbase import LangBB from devito.symbolics import c_complex, c_double_complex from devito.tools import dtype_to_cstr -from devito.petsc.utils import petsc_type_mappings + +from devito.petsc.config import petsc_type_mappings __all__ = ['CBB', 'CDataManager', 'COrchestrator'] @@ -82,3 +83,6 @@ class PetscCPrinter(CPrinter): _restrict_keyword = '' type_mappings = {**CPrinter.type_mappings, **petsc_type_mappings} + + def _print_Pi(self, expr): + return 'PETSC_PI' diff --git a/devito/petsc/config.py b/devito/petsc/config.py new file mode 100644 index 0000000000..e2fe3ed443 --- /dev/null +++ b/devito/petsc/config.py @@ -0,0 +1,86 @@ +import os +import ctypes +from pathlib import Path + +from petsctools import get_petscvariables, MissingPetscException + +from devito.tools import memoized_func + + +class PetscOSError(OSError): + pass + + +@memoized_func +def get_petsc_dir(): + petsc_dir = os.environ.get('PETSC_DIR') + if petsc_dir is None: + raise PetscOSError("PETSC_DIR environment variable not set") + else: + petsc_dir = (Path(petsc_dir),) + + petsc_arch = os.environ.get('PETSC_ARCH') + if petsc_arch is not None: + petsc_dir += (petsc_dir[0] / petsc_arch,) + + petsc_installed = petsc_dir[-1] / 'include' / 'petscconf.h' + if not petsc_installed.is_file(): + raise PetscOSError("PETSc is not installed") + + return petsc_dir + + +@memoized_func +def core_metadata(): + petsc_dir = get_petsc_dir() + + petsc_include = tuple([arch / 'include' for arch in petsc_dir]) + petsc_lib = tuple([arch / 'lib' for arch in petsc_dir]) + + return { + 'includes': ('petscsnes.h', 'petscdmda.h'), + 'include_dirs': petsc_include, + 'libs': ('petsc'), + 'lib_dirs': petsc_lib, + 'ldflags': tuple([f"-Wl,-rpath,{lib}" for lib in petsc_lib]) + } + + +try: + petsc_variables = get_petscvariables() +except MissingPetscException: + petsc_variables = {} + + +def get_petsc_type_mappings(): + try: + petsc_precision = petsc_variables['PETSC_PRECISION'] + except KeyError: + printer_mapper = {} + petsc_type_to_ctype = {} + else: + petsc_scalar = 'PetscScalar' + # TODO: Check to see whether Petsc is compiled with + # 32-bit or 64-bit integers + printer_mapper = {ctypes.c_int: 'PetscInt'} + + if petsc_precision == 'single': + printer_mapper[ctypes.c_float] = petsc_scalar + elif petsc_precision == 'double': + printer_mapper[ctypes.c_double] = petsc_scalar + + # Used to construct ctypes.Structures that wrap PETSc objects + petsc_type_to_ctype = {v: k for k, v in printer_mapper.items()} + # Add other PETSc types + petsc_type_to_ctype.update({ + 'KSPType': ctypes.c_char_p, + 'KSPConvergedReason': petsc_type_to_ctype['PetscInt'], + 'KSPNormType': petsc_type_to_ctype['PetscInt'], + }) + return printer_mapper, petsc_type_to_ctype + + +petsc_type_mappings, petsc_type_to_ctype = get_petsc_type_mappings() + + +petsc_languages = ['petsc'] diff --git a/devito/petsc/iet/builder.py b/devito/petsc/iet/builder.py new file mode 100644 index 0000000000..e1c178c059 --- /dev/null +++ b/devito/petsc/iet/builder.py @@ -0,0 +1,341 @@ +import math + +from devito.ir.iet import DummyExpr, BlankLine +from devito.symbolics import (Byref, FieldFromPointer, VOID, + FieldFromComposite, Null) + +from devito.petsc.iet.nodes import ( + FormFunctionCallback, MatShellSetOp, PETScCall, petsc_call +) + + +def make_core_petsc_calls(objs, comm): + call_mpi = petsc_call_mpi('MPI_Comm_size', [comm, Byref(objs['size'])]) + return call_mpi, BlankLine + + +class BuilderBase: + def __init__(self, **kwargs): + self.inject_solve = kwargs.get('inject_solve') + self.objs = kwargs.get('objs') + self.solver_objs = kwargs.get('solver_objs') + self.callback_builder = kwargs.get('callback_builder') + self.field_data = self.inject_solve.expr.rhs.field_data + self.formatted_prefix = self.inject_solve.expr.rhs.formatted_prefix + self.calls = self._setup() + + @property + def snes_ctx(self): + """ + The [optional] context for private data for the function evaluation routine. + https://petsc.org/main/manualpages/SNES/SNESSetFunction/ + """ + return VOID(self.solver_objs['dmda'], stars='*') + + def _setup(self): + sobjs = self.solver_objs + dmda = sobjs['dmda'] + + snes_create = petsc_call('SNESCreate', [sobjs['comm'], Byref(sobjs['snes'])]) + + snes_options_prefix = petsc_call( + 'SNESSetOptionsPrefix', [sobjs['snes'], sobjs['snes_prefix']] + ) if self.formatted_prefix else None + + set_options = petsc_call( + self.callback_builder._set_options_efunc.name, [] + ) + + snes_set_dm = petsc_call('SNESSetDM', [sobjs['snes'], dmda]) + + create_matrix = petsc_call('DMCreateMatrix', [dmda, Byref(sobjs['Jac'])]) + + snes_set_jac = petsc_call( + 'SNESSetJacobian', [sobjs['snes'], sobjs['Jac'], + sobjs['Jac'], 'MatMFFDComputeJacobian', Null] + ) + + global_x = petsc_call('DMCreateGlobalVector', + [dmda, Byref(sobjs['xglobal'])]) + + target = self.field_data.target + field_from_ptr = FieldFromPointer( + target.function._C_field_data, target.function._C_symbol + ) + + local_size = math.prod( + v for v, dim in zip(target.shape_allocated, target.dimensions) if dim.is_Space + ) + # TODO: Check - VecCreateSeqWithArray + local_x = petsc_call('VecCreateMPIWithArray', + [sobjs['comm'], 1, local_size, 'PETSC_DECIDE', + field_from_ptr, Byref(sobjs['xlocal'])]) + + # TODO: potentially also need to set the DM and local/global map to xlocal + + get_local_size = petsc_call('VecGetSize', + [sobjs['xlocal'], Byref(sobjs['localsize'])]) + + global_b = petsc_call('DMCreateGlobalVector', + [dmda, Byref(sobjs['bglobal'])]) + + snes_get_ksp = petsc_call('SNESGetKSP', + [sobjs['snes'], Byref(sobjs['ksp'])]) + + matvec = self.callback_builder.main_matvec_callback + matvec_operation = petsc_call( + 'MatShellSetOperation', + [sobjs['Jac'], 'MATOP_MULT', MatShellSetOp(matvec.name, void, void)] + ) + formfunc = self.callback_builder._F_efunc + formfunc_operation = petsc_call( + 'SNESSetFunction', + [sobjs['snes'], Null, FormFunctionCallback(formfunc.name, void, void), + self.snes_ctx] + ) + + snes_set_options = petsc_call( + 'SNESSetFromOptions', [sobjs['snes']] + ) + + dmda_calls = self._create_dmda_calls(dmda) + + mainctx = sobjs['userctx'] + + call_struct_callback = petsc_call( + self.callback_builder.user_struct_callback.name, [Byref(mainctx)] + ) + + # TODO: maybe don't need to explictly set this + mat_set_dm = petsc_call('MatSetDM', [sobjs['Jac'], dmda]) + + calls_set_app_ctx = petsc_call('DMSetApplicationContext', [dmda, Byref(mainctx)]) + + base_setup = dmda_calls + ( + snes_create, + snes_options_prefix, + set_options, + snes_set_dm, + create_matrix, + snes_set_jac, + global_x, + local_x, + get_local_size, + global_b, + snes_get_ksp, + matvec_operation, + formfunc_operation, + snes_set_options, + call_struct_callback, + mat_set_dm, + calls_set_app_ctx, + BlankLine + ) + extended_setup = self._extend_setup() + return base_setup + extended_setup + + def _extend_setup(self): + """ + Hook for subclasses to add additional setup calls. + """ + return () + + def _create_dmda_calls(self, dmda): + dmda_create = self._create_dmda(dmda) + dm_setup = petsc_call('DMSetUp', [dmda]) + dm_mat_type = petsc_call('DMSetMatType', [dmda, 'MATSHELL']) + return dmda_create, dm_setup, dm_mat_type + + def _create_dmda(self, dmda): + sobjs = self.solver_objs + grid = self.field_data.grid + nspace_dims = len(grid.dimensions) + + # MPI communicator + args = [sobjs['comm']] + + # Type of ghost nodes + args.extend(['DM_BOUNDARY_GHOSTED' for _ in range(nspace_dims)]) + + # Stencil type + if nspace_dims > 1: + args.append('DMDA_STENCIL_BOX') + + # Global dimensions + args.extend(list(grid.shape)[::-1]) + # No.of processors in each dimension + if nspace_dims > 1: + args.extend(list(grid.distributor.topology)[::-1]) + + # Number of degrees of freedom per node + args.append(dmda.dofs) + # "Stencil width" -> size of overlap + # TODO: Instead, this probably should be + # extracted from field_data.target._size_outhalo? + stencil_width = self.field_data.space_order + + args.append(stencil_width) + args.extend([Null]*nspace_dims) + + # The distributed array object + args.append(Byref(dmda)) + + # The PETSc call used to create the DMDA + dmda = petsc_call(f'DMDACreate{nspace_dims}d', args) + + return dmda + + +class CoupledBuilder(BuilderBase): + def _setup(self): + # TODO: minimise code duplication with superclass + objs = self.objs + sobjs = self.solver_objs + dmda = sobjs['dmda'] + + snes_create = petsc_call('SNESCreate', [sobjs['comm'], Byref(sobjs['snes'])]) + + snes_options_prefix = petsc_call( + 'SNESSetOptionsPrefix', [sobjs['snes'], sobjs['snes_prefix']] + ) if self.formatted_prefix else None + + set_options = petsc_call( + self.callback_builder._set_options_efunc.name, [] + ) + + snes_set_dm = petsc_call('SNESSetDM', [sobjs['snes'], dmda]) + + create_matrix = petsc_call('DMCreateMatrix', [dmda, Byref(sobjs['Jac'])]) + + snes_set_jac = petsc_call( + 'SNESSetJacobian', [sobjs['snes'], sobjs['Jac'], + sobjs['Jac'], 'MatMFFDComputeJacobian', Null] + ) + + global_x = petsc_call('DMCreateGlobalVector', + [dmda, Byref(sobjs['xglobal'])]) + + local_x = petsc_call('DMCreateLocalVector', [dmda, Byref(sobjs['xlocal'])]) + + get_local_size = petsc_call('VecGetSize', + [sobjs['xlocal'], Byref(sobjs['localsize'])]) + + snes_get_ksp = petsc_call('SNESGetKSP', + [sobjs['snes'], Byref(sobjs['ksp'])]) + + matvec = self.callback_builder.main_matvec_callback + matvec_operation = petsc_call( + 'MatShellSetOperation', + [sobjs['Jac'], 'MATOP_MULT', MatShellSetOp(matvec.name, void, void)] + ) + formfunc = self.callback_builder._F_efunc + formfunc_operation = petsc_call( + 'SNESSetFunction', + [sobjs['snes'], Null, FormFunctionCallback(formfunc.name, void, void), + self.snes_ctx] + ) + + snes_set_options = petsc_call( + 'SNESSetFromOptions', [sobjs['snes']] + ) + + dmda_calls = self._create_dmda_calls(dmda) + + mainctx = sobjs['userctx'] + + call_struct_callback = petsc_call( + self.callback_builder.user_struct_callback.name, [Byref(mainctx)] + ) + + # TODO: maybe don't need to explictly set this + mat_set_dm = petsc_call('MatSetDM', [sobjs['Jac'], dmda]) + + calls_set_app_ctx = petsc_call('DMSetApplicationContext', [dmda, Byref(mainctx)]) + + create_field_decomp = petsc_call( + 'DMCreateFieldDecomposition', + [dmda, Byref(sobjs['nfields']), Null, Byref(sobjs['fields']), + Byref(sobjs['subdms'])] + ) + submat_cb = self.callback_builder.submatrices_callback + matop_create_submats_op = petsc_call( + 'MatShellSetOperation', + [sobjs['Jac'], 'MATOP_CREATE_SUBMATRICES', + MatShellSetOp(submat_cb.name, void, void)] + ) + + call_coupled_struct_callback = petsc_call( + 'PopulateMatContext', + [Byref(sobjs['jacctx']), sobjs['subdms'], sobjs['fields']] + ) + + shell_set_ctx = petsc_call( + 'MatShellSetContext', [sobjs['Jac'], Byref(sobjs['jacctx']._C_symbol)] + ) + + create_submats = petsc_call( + 'MatCreateSubMatrices', + [sobjs['Jac'], sobjs['nfields'], sobjs['fields'], + sobjs['fields'], 'MAT_INITIAL_MATRIX', + Byref(FieldFromComposite(objs['Submats'].base, sobjs['jacctx']))] + ) + + targets = self.field_data.targets + + deref_dms = [ + DummyExpr(sobjs[f'da{t.name}'], sobjs['subdms'].indexed[i]) + for i, t in enumerate(targets) + ] + + xglobals = [petsc_call( + 'DMCreateGlobalVector', + [sobjs[f'da{t.name}'], Byref(sobjs[f'xglobal{t.name}'])] + ) for t in targets] + + xlocals = [] + for t in targets: + target_xloc = sobjs[f'xlocal{t.name}'] + local_size = math.prod( + v for v, dim in zip(t.shape_allocated, t.dimensions) if dim.is_Space + ) + field_from_ptr = FieldFromPointer( + t.function._C_field_data, t.function._C_symbol + ) + # TODO: Check - VecCreateSeqWithArray? + xlocals.append(petsc_call( + 'VecCreateMPIWithArray', + [sobjs['comm'], 1, local_size, 'PETSC_DECIDE', + field_from_ptr, Byref(target_xloc)] + )) + + coupled_setup = dmda_calls + ( + snes_create, + snes_options_prefix, + set_options, + snes_set_dm, + create_matrix, + snes_set_jac, + global_x, + local_x, + get_local_size, + snes_get_ksp, + matvec_operation, + formfunc_operation, + snes_set_options, + call_struct_callback, + mat_set_dm, + calls_set_app_ctx, + create_field_decomp, + matop_create_submats_op, + call_coupled_struct_callback, + shell_set_ctx, + create_submats) + \ + tuple(deref_dms) + tuple(xglobals) + tuple(xlocals) + (BlankLine,) + return coupled_setup + + +def petsc_call_mpi(specific_call, call_args): + return PETScCall('PetscCallMPI', [PETScCall(specific_call, arguments=call_args)]) + + +void = VOID._dtype diff --git a/devito/petsc/iet/callbacks.py b/devito/petsc/iet/callbacks.py new file mode 100644 index 0000000000..4d06063242 --- /dev/null +++ b/devito/petsc/iet/callbacks.py @@ -0,0 +1,1132 @@ +from collections import OrderedDict + +from devito.ir.iet import ( + Call, FindSymbols, List, Uxreplace, CallableBody, Dereference, DummyExpr, + BlankLine, Callable, Iteration, PointerCast, Definition +) +from devito.symbolics import ( + Byref, FieldFromPointer, IntDiv, Deref, Mod, String, Null, VOID +) +from devito.symbolics.unevaluation import Mul +from devito.types.basic import AbstractFunction +from devito.types import Dimension, Temp, TempArray +from devito.tools import filter_ordered + +from devito.petsc.iet.nodes import PETScCallable, MatShellSetOp, petsc_call +from devito.petsc.types import DMCast, MainUserStruct, CallbackUserStruct +from devito.petsc.iet.type_builder import objs +from devito.petsc.types.macros import petsc_func_begin_user +from devito.petsc.types.modes import InsertMode + + +class BaseCallbackBuilder: + """ + Build IET routines to generate PETSc callback functions. + """ + def __init__(self, **kwargs): + + self.rcompile = kwargs.get('rcompile', None) + self.sregistry = kwargs.get('sregistry', None) + self.concretize_mapper = kwargs.get('concretize_mapper', {}) + self.time_dependence = kwargs.get('time_dependence') + self.objs = kwargs.get('objs') + self.solver_objs = kwargs.get('solver_objs') + self.inject_solve = kwargs.get('inject_solve') + self.solve_expr = self.inject_solve.expr.rhs + + self._efuncs = OrderedDict() + self._struct_params = [] + + self._set_options_efunc = None + self._clear_options_efunc = None + self._main_matvec_callback = None + self._user_struct_callback = None + self._F_efunc = None + self._b_efunc = None + + self._J_efuncs = [] + self._initial_guesses = [] + + self._make_core() + self._efuncs = self._uxreplace_efuncs() + + @property + def efuncs(self): + return self._efuncs + + @property + def struct_params(self): + return self._struct_params + + @property + def filtered_struct_params(self): + return filter_ordered(self.struct_params) + + @property + def main_matvec_callback(self): + """ + The matrix-vector callback for the full Jacobian. + This is the function set in the main Kernel via: + PetscCall(MatShellSetOperation(J, MATOP_MULT, (void (*)(void))...)); + The callback has the signature `(Mat, Vec, Vec)`. + """ + return self._J_efuncs[0] + + @property + def J_efuncs(self): + """ + List of matrix-vector callbacks. + Each callback has the signature `(Mat, Vec, Vec)`. Typically, this list + contains a single element, but in mixed systems it can include multiple + callbacks, one for each subblock. + """ + return self._J_efuncs + + @property + def initial_guesses(self): + return self._initial_guesses + + @property + def user_struct_callback(self): + return self._user_struct_callback + + @property + def solver_parameters(self): + return self.solve_expr.solver_parameters + + @property + def field_data(self): + return self.solve_expr.field_data + + @property + def formatted_prefix(self): + return self.solve_expr.formatted_prefix + + @property + def arrays(self): + return self.field_data.arrays + + @property + def target(self): + return self.field_data.target + + def _make_core(self): + self._make_options_callback() + self._make_matvec(self.field_data.jacobian) + self._make_formfunc() + self._make_formrhs() + if self.field_data.initial_guess.exprs: + self._make_initial_guess() + self._make_user_struct_callback() + + def _make_petsc_callable(self, prefix, body, parameters=()): + return PETScCallable( + self.sregistry.make_name(prefix=prefix), + body, + retval=self.objs['err'], + parameters=parameters + ) + + def _make_callable_body(self, body, standalones=(), stacks=(), casts=()): + return CallableBody( + List(body=body), + init=(petsc_func_begin_user,), + standalones=standalones, + stacks=stacks, + casts=casts, + retstmt=(Call('PetscFunctionReturn', arguments=[0]),) + ) + + def _make_options_callback(self): + """ + Create two callbacks: one to set PETSc options and one + to clear them. + Options are only set/cleared if they were not specifed via + command line arguments. + """ + params = self.solver_parameters + prefix = self.inject_solve.expr.rhs.formatted_prefix + + set_body, clear_body = [], [] + + for k, v in params.items(): + option = f'-{prefix}{k}' + + # TODO: Revisit use of a global variable here. + # Consider replacing this with a call to `PetscGetArgs`, though + # initial attempts failed, possibly because the argv pointer is + # created in Python?.. + import devito.petsc.initialize + if option in devito.petsc.initialize._petsc_clargs: + # Ensures that the command line args take priority + continue + + option_name = String(option) + # For options without a value e.g `ksp_view`, pass Null + option_value = Null if v is None else String(str(v)) + set_body.append( + petsc_call('PetscOptionsSetValue', [Null, option_name, option_value]) + ) + clear_body.append( + petsc_call('PetscOptionsClearValue', [Null, option_name]) + ) + + set_body = self._make_callable_body(set_body) + clear_body = self._make_callable_body(clear_body) + + set_callback = self._make_petsc_callable('SetPetscOptions', set_body) + clear_callback = self._make_petsc_callable('ClearPetscOptions', clear_body) + + self._set_options_efunc = set_callback + self._efuncs[set_callback.name] = set_callback + self._clear_options_efunc = clear_callback + self._efuncs[clear_callback.name] = clear_callback + + def _make_matvec(self, jacobian, prefix='MatMult'): + # Compile `matvecs` into an IET via recursive compilation + matvecs = jacobian.matvecs + irs, _ = self.rcompile( + matvecs, options={'mpi': False}, sregistry=self.sregistry, + concretize_mapper=self.concretize_mapper + ) + body = self._create_matvec_body( + List(body=irs.uiet.body), jacobian + ) + objs = self.objs + cb = self._make_petsc_callable( + prefix, body, parameters=(objs['J'], objs['X'], objs['Y']) + ) + self._J_efuncs.append(cb) + self._efuncs[cb.name] = cb + + def _create_matvec_body(self, body, jacobian): + linsolve_expr = self.inject_solve.expr.rhs + objs = self.objs + sobjs = self.solver_objs + + dmda = sobjs['callbackdm'] + ctx = objs['dummyctx'] + xlocal = objs['xloc'] + ylocal = objs['yloc'] + y_matvec = self.arrays[jacobian.row_target]['y'] + x_matvec = self.arrays[jacobian.col_target]['x'] + + body = self.time_dependence.uxreplace_time(body) + + fields = get_user_struct_fields(body) + + mat_get_dm = petsc_call('MatGetDM', [objs['J'], Byref(dmda)]) + + dm_get_app_context = petsc_call( + 'DMGetApplicationContext', [dmda, Byref(ctx._C_symbol)] + ) + + zero_y_memory = zero_vector(objs['Y']) if jacobian.zero_memory else None + + dm_get_local_xvec = petsc_call( + 'DMGetLocalVector', [dmda, Byref(xlocal)] + ) + + global_to_local_begin = petsc_call( + 'DMGlobalToLocalBegin', [dmda, objs['X'], insert_values, xlocal] + ) + + global_to_local_end = petsc_call('DMGlobalToLocalEnd', [ + dmda, objs['X'], insert_values, xlocal + ]) + + dm_get_local_yvec = petsc_call( + 'DMGetLocalVector', [dmda, Byref(ylocal)] + ) + + zero_ylocal_memory = zero_vector(ylocal) + + vec_get_array_y = petsc_call( + 'VecGetArray', [ylocal, Byref(y_matvec._C_symbol)] + ) + + vec_get_array_x = petsc_call( + 'VecGetArray', [xlocal, Byref(x_matvec._C_symbol)] + ) + + dm_get_local_info = petsc_call( + 'DMDAGetLocalInfo', [dmda, Byref(linsolve_expr.localinfo)] + ) + + vec_restore_array_y = petsc_call( + 'VecRestoreArray', [ylocal, Byref(y_matvec._C_symbol)] + ) + + vec_restore_array_x = petsc_call( + 'VecRestoreArray', [xlocal, Byref(x_matvec._C_symbol)] + ) + + dm_local_to_global_begin = petsc_call('DMLocalToGlobalBegin', [ + dmda, ylocal, add_values, objs['Y'] + ]) + + dm_local_to_global_end = petsc_call('DMLocalToGlobalEnd', [ + dmda, ylocal, add_values, objs['Y'] + ]) + + dm_restore_local_xvec = petsc_call( + 'DMRestoreLocalVector', [dmda, Byref(xlocal)] + ) + + dm_restore_local_yvec = petsc_call( + 'DMRestoreLocalVector', [dmda, Byref(ylocal)] + ) + + # 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, + dm_restore_local_xvec, + dm_restore_local_yvec) + ) + + stacks = ( + zero_y_memory, + dm_get_local_xvec, + global_to_local_begin, + global_to_local_end, + dm_get_local_yvec, + zero_ylocal_memory, + vec_get_array_y, + vec_get_array_x, + dm_get_local_info + ) + + # Dereference function data in struct + derefs = dereference_funcs(ctx, fields) + + # Force the struct definition to appear at the very start, since + # stacks, allocs etc may rely on its information + struct_definition = [ + Definition(ctx), Definition(dmda), mat_get_dm, dm_get_app_context + ] + + body = self._make_callable_body( + body, standalones=struct_definition, stacks=stacks+derefs + ) + + # Replace non-function data with pointer to data in struct + subs = {i._C_symbol: FieldFromPointer(i._C_symbol, ctx) for i in fields} + body = Uxreplace(subs).visit(body) + + self._struct_params.extend(fields) + return body + + def _make_formfunc(self): + objs = self.objs + F_exprs = self.field_data.residual.F_exprs + # Compile `F_exprs` into an IET via recursive compilation + irs, _ = self.rcompile( + F_exprs, options={'mpi': False}, sregistry=self.sregistry, + concretize_mapper=self.concretize_mapper + ) + body_formfunc = self._create_formfunc_body( + List(body=irs.uiet.body) + ) + parameters = (objs['snes'], objs['X'], objs['F'], objs['dummyptr']) + cb = self._make_petsc_callable('FormFunction', body_formfunc, parameters) + + self._F_efunc = cb + self._efuncs[cb.name] = cb + + def _create_formfunc_body(self, body): + linsolve_expr = self.inject_solve.expr.rhs + objs = self.objs + sobjs = self.solver_objs + arrays = self.arrays + target = self.target + + dmda = sobjs['callbackdm'] + ctx = objs['dummyctx'] + + body = self.time_dependence.uxreplace_time(body) + + fields = get_user_struct_fields(body) + self._struct_params.extend(fields) + + f_formfunc = arrays[target]['f'] + x_formfunc = arrays[target]['x'] + + dm_cast = DummyExpr(dmda, DMCast(objs['dummyptr']), init=True) + + dm_get_app_context = petsc_call( + 'DMGetApplicationContext', [dmda, Byref(ctx._C_symbol)] + ) + + zero_f_memory = zero_vector(objs['F']) + + dm_get_local_xvec = petsc_call( + 'DMGetLocalVector', [dmda, Byref(objs['xloc'])] + ) + + global_to_local_begin = petsc_call( + 'DMGlobalToLocalBegin', [dmda, objs['X'], insert_values, objs['xloc']] + ) + + global_to_local_end = petsc_call( + 'DMGlobalToLocalEnd', [dmda, objs['X'], insert_values, objs['xloc']] + ) + + dm_get_local_yvec = petsc_call( + 'DMGetLocalVector', [dmda, Byref(objs['floc'])] + ) + + vec_get_array_y = petsc_call( + 'VecGetArray', [objs['floc'], Byref(f_formfunc._C_symbol)] + ) + + vec_get_array_x = petsc_call( + 'VecGetArray', [objs['xloc'], Byref(x_formfunc._C_symbol)] + ) + + dm_get_local_info = petsc_call( + 'DMDAGetLocalInfo', [dmda, Byref(linsolve_expr.localinfo)] + ) + + vec_restore_array_y = petsc_call( + 'VecRestoreArray', [objs['floc'], Byref(f_formfunc._C_symbol)] + ) + + vec_restore_array_x = petsc_call( + 'VecRestoreArray', [objs['xloc'], Byref(x_formfunc._C_symbol)] + ) + + dm_local_to_global_begin = petsc_call('DMLocalToGlobalBegin', [ + dmda, objs['floc'], add_values, objs['F'] + ]) + + dm_local_to_global_end = petsc_call('DMLocalToGlobalEnd', [ + dmda, objs['floc'], add_values, objs['F'] + ]) + + dm_restore_local_xvec = petsc_call( + 'DMRestoreLocalVector', [dmda, Byref(objs['xloc'])] + ) + + dm_restore_local_yvec = petsc_call( + 'DMRestoreLocalVector', [dmda, Byref(objs['floc'])] + ) + + body = body._rebuild( + body=body.body + + (vec_restore_array_y, + vec_restore_array_x, + dm_local_to_global_begin, + dm_local_to_global_end, + dm_restore_local_xvec, + dm_restore_local_yvec) + ) + + stacks = ( + zero_f_memory, + 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 + ) + + # Dereference function data in struct + derefs = dereference_funcs(ctx, fields) + + # Force the struct definition to appear at the very start, since + # stacks, allocs etc may rely on its information + struct_definition = [Definition(ctx), dm_cast, dm_get_app_context] + + body = self._make_callable_body( + body, standalones=struct_definition, stacks=stacks+derefs + ) + # Replace non-function data with pointer to data in struct + subs = {i._C_symbol: FieldFromPointer(i._C_symbol, ctx) for i in fields} + + return Uxreplace(subs).visit(body) + + def _make_formrhs(self): + b_exprs = self.field_data.residual.b_exprs + sobjs = self.solver_objs + + # Compile `b_exprs` into an IET via recursive compilation + irs, _ = self.rcompile( + b_exprs, options={'mpi': False}, sregistry=self.sregistry, + concretize_mapper=self.concretize_mapper + ) + body = self._create_form_rhs_body( + List(body=irs.uiet.body) + ) + objs = self.objs + cb = self._make_petsc_callable( + 'FormRHS', body, parameters=(sobjs['callbackdm'], objs['B']) + ) + self._b_efunc = cb + self._efuncs[cb.name] = cb + + def _create_form_rhs_body(self, body): + linsolve_expr = self.inject_solve.expr.rhs + objs = self.objs + sobjs = self.solver_objs + target = self.target + + dmda = sobjs['callbackdm'] + ctx = objs['dummyctx'] + + dm_get_local = petsc_call( + 'DMGetLocalVector', [dmda, Byref(sobjs['blocal'])] + ) + + dm_global_to_local_begin = petsc_call( + 'DMGlobalToLocalBegin', [dmda, objs['B'], insert_values, sobjs['blocal']] + ) + + dm_global_to_local_end = petsc_call( + 'DMGlobalToLocalEnd', [dmda, objs['B'], insert_values, sobjs['blocal']] + ) + + b_arr = self.field_data.arrays[target]['b'] + + vec_get_array = petsc_call( + 'VecGetArray', [sobjs['blocal'], Byref(b_arr._C_symbol)] + ) + + dm_get_local_info = petsc_call( + 'DMDAGetLocalInfo', [dmda, Byref(linsolve_expr.localinfo)] + ) + + body = self.time_dependence.uxreplace_time(body) + + fields = get_user_struct_fields(body) + self._struct_params.extend(fields) + + dm_get_app_context = petsc_call( + 'DMGetApplicationContext', [dmda, Byref(ctx._C_symbol)] + ) + + dm_local_to_global_begin = petsc_call('DMLocalToGlobalBegin', [ + dmda, sobjs['blocal'], insert_values, objs['B'] + ]) + + dm_local_to_global_end = petsc_call('DMLocalToGlobalEnd', [ + dmda, sobjs['blocal'], insert_values, objs['B'] + ]) + + vec_restore_array = petsc_call( + 'VecRestoreArray', [sobjs['blocal'], Byref(b_arr._C_symbol)] + ) + + dm_restore_local_bvec = petsc_call( + 'DMRestoreLocalVector', [dmda, Byref(sobjs['blocal'])] + ) + + body = body._rebuild(body=body.body + ( + dm_local_to_global_begin, dm_local_to_global_end, vec_restore_array, + dm_restore_local_bvec + )) + + stacks = ( + dm_get_local, + dm_global_to_local_begin, + dm_global_to_local_end, + vec_get_array, + dm_get_local_info + ) + + # Dereference function data in struct + derefs = dereference_funcs(ctx, fields) + + # Force the struct definition to appear at the very start, since + # stacks, allocs etc may rely on its information + struct_definition = [Definition(ctx), dm_get_app_context] + + body = self._make_callable_body( + [body], standalones=struct_definition, stacks=stacks+derefs + ) + + # Replace non-function data with pointer to data in struct + subs = {i._C_symbol: FieldFromPointer(i._C_symbol, ctx) for + i in fields if not isinstance(i.function, AbstractFunction)} + + return Uxreplace(subs).visit(body) + + def _make_initial_guess(self): + exprs = self.field_data.initial_guess.exprs + sobjs = self.solver_objs + objs = self.objs + + # Compile initital guess `eqns` into an IET via recursive compilation + irs, _ = self.rcompile( + exprs, options={'mpi': False}, sregistry=self.sregistry, + concretize_mapper=self.concretize_mapper + ) + body = self._create_initial_guess_body( + List(body=irs.uiet.body) + ) + cb = self._make_petsc_callable( + 'FormInitialGuess', body, parameters=(sobjs['callbackdm'], objs['xloc']) + ) + self._initial_guesses.append(cb) + self._efuncs[cb.name] = cb + + def _create_initial_guess_body(self, body): + linsolve_expr = self.inject_solve.expr.rhs + objs = self.objs + sobjs = self.solver_objs + target = self.target + + dmda = sobjs['callbackdm'] + ctx = objs['dummyctx'] + + x_arr = self.field_data.arrays[target]['x'] + + vec_get_array = petsc_call( + 'VecGetArray', [objs['xloc'], Byref(x_arr._C_symbol)] + ) + + dm_get_local_info = petsc_call( + 'DMDAGetLocalInfo', [dmda, Byref(linsolve_expr.localinfo)] + ) + + body = self.time_dependence.uxreplace_time(body) + + fields = get_user_struct_fields(body) + self._struct_params.extend(fields) + + dm_get_app_context = petsc_call( + 'DMGetApplicationContext', [dmda, Byref(ctx._C_symbol)] + ) + + vec_restore_array = petsc_call( + 'VecRestoreArray', [objs['xloc'], Byref(x_arr._C_symbol)] + ) + + body = body._rebuild(body=body.body + (vec_restore_array,)) + + stacks = ( + vec_get_array, + dm_get_local_info + ) + + # Dereference function data in struct + derefs = dereference_funcs(ctx, fields) + + # Force the struct definition to appear at the very start, since + # stacks, allocs etc may rely on its information + struct_definition = [Definition(ctx), dm_get_app_context] + + body = self._make_callable_body( + body, standalones=struct_definition, stacks=stacks+derefs + ) + + # Replace non-function data with pointer to data in struct + subs = {i._C_symbol: FieldFromPointer(i._C_symbol, ctx) for + i in fields if not isinstance(i.function, AbstractFunction)} + + return Uxreplace(subs).visit(body) + + def _make_user_struct_callback(self): + """ + This is the struct initialised inside the main kernel and + attached to the DM via DMSetApplicationContext. + """ + mainctx = self.solver_objs['userctx'] = MainUserStruct( + name=self.sregistry.make_name(prefix='ctx'), + pname=self.sregistry.make_name(prefix='UserCtx'), + fields=self.filtered_struct_params, + liveness='lazy', + modifier=None + ) + body = [ + DummyExpr(FieldFromPointer(i._C_symbol, mainctx), i._C_symbol) + for i in mainctx.callback_fields + ] + struct_callback_body = self._make_callable_body(body) + cb = Callable( + self.sregistry.make_name(prefix='PopulateUserContext'), + struct_callback_body, self.objs['err'], + parameters=[mainctx] + ) + self._efuncs[cb.name] = cb + self._user_struct_callback = cb + + def _uxreplace_efuncs(self): + sobjs = self.solver_objs + callback_user_struct = CallbackUserStruct( + name=sobjs['userctx'].name, + pname=sobjs['userctx'].pname, + fields=self.filtered_struct_params, + liveness='lazy', + modifier=' *', + parent=sobjs['userctx'] + ) + mapper = {} + visitor = Uxreplace({self.objs['dummyctx']: callback_user_struct}) + for k, v in self._efuncs.items(): + mapper.update({k: visitor.visit(v)}) + return mapper + + +class CoupledCallbackBuilder(BaseCallbackBuilder): + def __init__(self, **kwargs): + self._submatrices_callback = None + super().__init__(**kwargs) + + @property + def submatrices_callback(self): + return self._submatrices_callback + + @property + def jacobian(self): + return self.inject_solve.expr.rhs.field_data.jacobian + + @property + def main_matvec_callback(self): + """ + This is the matrix-vector callback associated with the whole Jacobian i.e + is set in the main kernel via + `PetscCall(MatShellSetOperation(J,MATOP_MULT,(void (*)(void))MyMatShellMult));` + """ + return self._main_matvec_callback + + def _make_core(self): + for sm in self.field_data.jacobian.nonzero_submatrices: + self._make_matvec(sm, prefix=f'{sm.name}_MatMult') + + self._make_options_callback() + self._make_whole_matvec() + self._make_whole_formfunc() + self._make_user_struct_callback() + self._create_submatrices() + self._efuncs['PopulateMatContext'] = self.objs['dummyefunc'] + + def _make_whole_matvec(self): + objs = self.objs + body = self._whole_matvec_body() + + parameters = (objs['J'], objs['X'], objs['Y']) + cb = self._make_petsc_callable( + 'WholeMatMult', List(body=body), parameters=parameters + ) + self._main_matvec_callback = cb + self._efuncs[cb.name] = cb + + def _whole_matvec_body(self): + objs = self.objs + sobjs = self.solver_objs + + jctx = objs['ljacctx'] + ctx_main = petsc_call('MatShellGetContext', [objs['J'], Byref(jctx)]) + + nonzero_submats = self.jacobian.nonzero_submatrices + + zero_y_memory = zero_vector(objs['Y']) + + calls = () + for sm in nonzero_submats: + name = sm.name + ctx = sobjs[f'{name}ctx'] + X = sobjs[f'{name}X'] + Y = sobjs[f'{name}Y'] + rows = objs['rows'].base + cols = objs['cols'].base + sm_indexed = objs['Submats'].indexed[sm.linear_idx] + + calls += ( + DummyExpr(sobjs[name], FieldFromPointer(sm_indexed, jctx)), + petsc_call('MatShellGetContext', [sobjs[name], Byref(ctx)]), + petsc_call( + 'VecGetSubVector', + [objs['X'], Deref(FieldFromPointer(cols, ctx)), Byref(X)] + ), + petsc_call( + 'VecGetSubVector', + [objs['Y'], Deref(FieldFromPointer(rows, ctx)), Byref(Y)] + ), + petsc_call('MatMult', [sobjs[name], X, Y]), + petsc_call( + 'VecRestoreSubVector', + [objs['X'], Deref(FieldFromPointer(cols, ctx)), Byref(X)] + ), + petsc_call( + 'VecRestoreSubVector', + [objs['Y'], Deref(FieldFromPointer(rows, ctx)), Byref(Y)] + ), + ) + body = (ctx_main, zero_y_memory, BlankLine) + calls + return self._make_callable_body(body) + + def _make_whole_formfunc(self): + objs = self.objs + F_exprs = self.field_data.residual.F_exprs + # Compile `F_exprs` into an IET via recursive compilation + irs, _ = self.rcompile( + F_exprs, options={'mpi': False}, sregistry=self.sregistry, + concretize_mapper=self.concretize_mapper + ) + body = self._whole_formfunc_body(List(body=irs.uiet.body)) + + parameters = (objs['snes'], objs['X'], objs['F'], objs['dummyptr']) + cb = self._make_petsc_callable( + 'WholeFormFunc', body, parameters=parameters + ) + + self._F_efunc = cb + self._efuncs[cb.name] = cb + + def _whole_formfunc_body(self, body): + linsolve_expr = self.inject_solve.expr.rhs + objs = self.objs + sobjs = self.solver_objs + + dmda = sobjs['callbackdm'] + ctx = objs['dummyctx'] + + body = self.time_dependence.uxreplace_time(body) + + fields = get_user_struct_fields(body) + self._struct_params.extend(fields) + + # Process body with bundles for residual callback + bundles = sobjs['bundles'] + fbundle = bundles['f'] + xbundle = bundles['x'] + + body = self.residual_bundle(body, bundles) + + dm_cast = DummyExpr(dmda, DMCast(objs['dummyptr']), init=True) + + dm_get_app_context = petsc_call( + 'DMGetApplicationContext', [dmda, Byref(ctx._C_symbol)] + ) + + zero_f_memory = zero_vector(objs['F']) + + dm_get_local_xvec = petsc_call( + 'DMGetLocalVector', [dmda, Byref(objs['xloc'])] + ) + + global_to_local_begin = petsc_call('DMGlobalToLocalBegin', [ + dmda, objs['X'], insert_values, objs['xloc'] + ]) + + global_to_local_end = petsc_call('DMGlobalToLocalEnd', [ + dmda, objs['X'], insert_values, objs['xloc'] + ]) + + dm_get_local_yvec = petsc_call( + 'DMGetLocalVector', [dmda, Byref(objs['floc'])] + ) + + vec_get_array_f = petsc_call( + 'VecGetArray', [objs['floc'], Byref(fbundle.vector._C_symbol)] + ) + + vec_get_array_x = petsc_call( + 'VecGetArray', [objs['xloc'], Byref(xbundle.vector._C_symbol)] + ) + + dm_get_local_info = petsc_call( + 'DMDAGetLocalInfo', [dmda, Byref(linsolve_expr.localinfo)] + ) + + vec_restore_array_f = petsc_call( + 'VecRestoreArray', [objs['floc'], Byref(fbundle.vector._C_symbol)] + ) + + vec_restore_array_x = petsc_call( + 'VecRestoreArray', [objs['xloc'], Byref(xbundle.vector._C_symbol)] + ) + + dm_local_to_global_begin = petsc_call('DMLocalToGlobalBegin', [ + dmda, objs['floc'], add_values, objs['F'] + ]) + + dm_local_to_global_end = petsc_call('DMLocalToGlobalEnd', [ + dmda, objs['floc'], add_values, objs['F'] + ]) + + dm_restore_local_xvec = petsc_call( + 'DMRestoreLocalVector', [dmda, Byref(objs['xloc'])] + ) + + dm_restore_local_yvec = petsc_call( + 'DMRestoreLocalVector', [dmda, Byref(objs['floc'])] + ) + + body = body._rebuild( + body=body.body + + (vec_restore_array_f, + vec_restore_array_x, + dm_local_to_global_begin, + dm_local_to_global_end, + dm_restore_local_xvec, + dm_restore_local_yvec) + ) + + stacks = ( + zero_f_memory, + dm_get_local_xvec, + global_to_local_begin, + global_to_local_end, + dm_get_local_yvec, + vec_get_array_f, + vec_get_array_x, + dm_get_local_info + ) + + # Dereference function data in struct + derefs = dereference_funcs(ctx, fields) + + # Force the struct definition to appear at the very start, since + # stacks, allocs etc may rely on its information + struct_definition = [Definition(ctx), dm_cast, dm_get_app_context] + + f_soa = PointerCast(fbundle) + x_soa = PointerCast(xbundle) + + formfunc_body = self._make_callable_body( + body, + standalones=struct_definition, + stacks=stacks+derefs, + casts=(f_soa, x_soa), + ) + # Replace non-function data with pointer to data in struct + subs = {i._C_symbol: FieldFromPointer(i._C_symbol, ctx) for i in fields} + + return Uxreplace(subs).visit(formfunc_body) + + def _create_submatrices(self): + body = self._submat_callback_body() + objs = self.objs + params = ( + objs['J'], + objs['nfields'], + objs['irow'], + objs['icol'], + objs['matreuse'], + objs['Submats'], + ) + cb = self._make_petsc_callable( + 'MatCreateSubMatrices', body, parameters=params) + + self._submatrices_callback = cb + self._efuncs[cb.name] = cb + + def _submat_callback_body(self): + objs = self.objs + sobjs = self.solver_objs + + n_submats = DummyExpr( + objs['nsubmats'], Mul(objs['nfields'], objs['nfields']) + ) + + malloc_submats = petsc_call( + 'PetscCalloc1', [objs['nsubmats'], objs['Submats']._C_symbol] + ) + + mat_get_dm = petsc_call('MatGetDM', [objs['J'], Byref(sobjs['callbackdm'])]) + + dm_get_app = petsc_call( + 'DMGetApplicationContext', [sobjs['callbackdm'], Byref(objs['dummyctx'])] + ) + + get_ctx = petsc_call('MatShellGetContext', [objs['J'], Byref(objs['ljacctx'])]) + + dm_get_info = petsc_call( + 'DMDAGetInfo', [ + sobjs['callbackdm'], Null, Byref(sobjs['M']), Byref(sobjs['N']), + Null, Null, Null, Null, Byref(objs['dof']), Null, Null, Null, Null, Null + ] + ) + subblock_rows = DummyExpr(objs['subblockrows'], Mul(sobjs['M'], sobjs['N'])) + subblock_cols = DummyExpr(objs['subblockcols'], Mul(sobjs['M'], sobjs['N'])) + + ptr = DummyExpr( + objs['submat_arr']._C_symbol, Deref(objs['Submats']._C_symbol), init=True + ) + + mat_create = petsc_call('MatCreate', [sobjs['comm'], Byref(objs['block'])]) + + mat_set_sizes = petsc_call( + 'MatSetSizes', [ + objs['block'], 'PETSC_DECIDE', 'PETSC_DECIDE', + objs['subblockrows'], objs['subblockcols'] + ] + ) + + mat_set_type = petsc_call('MatSetType', [objs['block'], 'MATSHELL']) + + malloc = petsc_call('PetscMalloc1', [1, Byref(objs['subctx'])]) + i = Dimension(name='i') + + row_idx = DummyExpr(objs['rowidx'], IntDiv(i, objs['dof'])) + col_idx = DummyExpr(objs['colidx'], Mod(i, objs['dof'])) + + deref_subdm = Dereference(objs['Subdms'], objs['ljacctx']) + + set_rows = DummyExpr( + FieldFromPointer(objs['rows'].base, objs['subctx']), + Byref(objs['irow'].indexed[objs['rowidx']]) + ) + set_cols = DummyExpr( + FieldFromPointer(objs['cols'].base, objs['subctx']), + Byref(objs['icol'].indexed[objs['colidx']]) + ) + dm_set_ctx = petsc_call( + 'DMSetApplicationContext', [ + objs['Subdms'].indexed[objs['rowidx']], objs['dummyctx'] + ] + ) + matset_dm = petsc_call('MatSetDM', [ + objs['block'], objs['Subdms'].indexed[objs['rowidx']] + ]) + + set_ctx = petsc_call('MatShellSetContext', [objs['block'], objs['subctx']]) + + mat_setup = petsc_call('MatSetUp', [objs['block']]) + + assign_block = DummyExpr(objs['submat_arr'].indexed[i], objs['block']) + + iter_body = ( + mat_create, + mat_set_sizes, + mat_set_type, + malloc, + row_idx, + col_idx, + set_rows, + set_cols, + dm_set_ctx, + matset_dm, + set_ctx, + mat_setup, + assign_block + ) + + upper_bound = objs['nsubmats'] - 1 + iteration = Iteration(List(body=iter_body), i, upper_bound) + + nonzero_submats = self.jacobian.nonzero_submatrices + matvec_lookup = {mv.name.split('_')[0]: mv for mv in self.J_efuncs} + + matmult_op = [ + petsc_call( + 'MatShellSetOperation', + [ + objs['submat_arr'].indexed[sb.linear_idx], + 'MATOP_MULT', + MatShellSetOp(matvec_lookup[sb.name].name, VOID._dtype, VOID._dtype), + ], + ) + for sb in nonzero_submats if sb.name in matvec_lookup + ] + + body = [ + n_submats, + malloc_submats, + mat_get_dm, + dm_get_app, + dm_get_info, + subblock_rows, + subblock_cols, + ptr, + BlankLine, + iteration, + ] + matmult_op + return self._make_callable_body(tuple(body), stacks=(get_ctx, deref_subdm)) + + def residual_bundle(self, body, bundles): + """ + Replaces PetscArrays in `body` with PetscBundle struct field accesses + (e.g., f_v[ix][iy] -> f_bundle[ix][iy].v). + Example: + f_v[ix][iy] = x_v[ix][iy]; + f_u[ix][iy] = x_u[ix][iy]; + becomes: + f_bundle[ix][iy].v = x_bundle[ix][iy].v; + f_bundle[ix][iy].u = x_bundle[ix][iy].u; + NOTE: This is used because the data is interleaved for + multi-component DMDAs in PETSc. + """ + mapper = bundles['bundle_mapper'] + indexeds = FindSymbols('indexeds').visit(body) + subs = {} + + for i in indexeds: + if i.base in mapper: + bundle = mapper[i.base] + index = bundles['target_indices'][i.function.target] + index = (index,) + i.indices + subs[i] = bundle.__getitem__(index) + + body = Uxreplace(subs).visit(body) + return body + + +def populate_matrix_context(efuncs): + if not objs['dummyefunc'] in efuncs.values(): + return + + subdms_expr = DummyExpr( + FieldFromPointer(objs['Subdms']._C_symbol, objs['ljacctx']), + objs['Subdms']._C_symbol + ) + fields_expr = DummyExpr( + FieldFromPointer(objs['Fields']._C_symbol, objs['ljacctx']), + objs['Fields']._C_symbol + ) + body = CallableBody( + List(body=[subdms_expr, fields_expr]), + init=(petsc_func_begin_user,), + retstmt=tuple([Call('PetscFunctionReturn', arguments=[0])]) + ) + name = 'PopulateMatContext' + efuncs[name] = Callable( + name, body, objs['err'], + parameters=[objs['ljacctx'], objs['Subdms'], objs['Fields']] + ) + + +def dereference_funcs(struct, fields): + """ + Dereference AbstractFunctions from a struct. + """ + return tuple( + [Dereference(i, struct) for i in + fields if isinstance(i.function, AbstractFunction)] + ) + + +def zero_vector(vec): + """ + Set all entries of a PETSc vector to zero. + """ + return petsc_call('VecSet', [vec, 0.0]) + + +def get_user_struct_fields(iet): + fields = [f.function for f in FindSymbols('basics').visit(iet)] + from devito.types.basic import LocalType + avoid = (Temp, TempArray, LocalType) + fields = [f for f in fields if not isinstance(f.function, avoid)] + fields = [ + f for f in fields if not (f.is_Dimension and not (f.is_Time or f.is_Modulo)) + ] + return fields + + +insert_values = InsertMode.insert_values +add_values = InsertMode.add_values diff --git a/devito/petsc/iet/logging.py b/devito/petsc/iet/logging.py index 68f31cfa91..65e2ec2be8 100644 --- a/devito/petsc/iet/logging.py +++ b/devito/petsc/iet/logging.py @@ -5,7 +5,7 @@ from devito.logger import PERF from devito.tools import frozendict -from devito.petsc.iet.utils import petsc_call +from devito.petsc.iet.nodes import petsc_call from devito.petsc.logging import petsc_return_variable_dict, PetscInfo diff --git a/devito/petsc/iet/nodes.py b/devito/petsc/iet/nodes.py index abb5da3acd..70508970c3 100644 --- a/devito/petsc/iet/nodes.py +++ b/devito/petsc/iet/nodes.py @@ -29,3 +29,12 @@ def callback_form(self): class PETScCall(Call): pass + + +def petsc_call(specific_call, call_args): + return PETScCall('PetscCall', [PETScCall(specific_call, arguments=call_args)]) + + +# Mapping special Eq operations to their corresponding IET Expression subclass types. +# These operations correspond to subclasses of `Eq`` utilised within `petscsolve``. +petsc_iet_mapper = {OpPetsc: PetscMetaData} diff --git a/devito/petsc/iet/passes.py b/devito/petsc/iet/passes.py index eb6cf48a1f..5154afe43d 100644 --- a/devito/petsc/iet/passes.py +++ b/devito/petsc/iet/passes.py @@ -3,34 +3,35 @@ from functools import cached_property from devito.passes.iet.engine import iet_pass -from devito.ir.iet import (Transformer, MapNodes, Iteration, BlankLine, - DummyExpr, CallableBody, List, Call, Callable, - FindNodes, Section) -from devito.symbolics import Byref, FieldFromPointer, Macro, Null -from devito.types import Symbol, Scalar +from devito.ir.iet import ( + Transformer, MapNodes, Iteration, CallableBody, List, Call, FindNodes, Section, + FindSymbols, DummyExpr, Uxreplace, Dereference +) +from devito.symbolics import Byref, Macro, Null, FieldFromPointer from devito.types.basic import DataSymbol -from devito.tools import frozendict import devito.logger -from devito.petsc.types import (PetscMPIInt, PetscErrorCode, MultipleFieldData, - PointerIS, Mat, CallbackVec, Vec, CallbackMat, SNES, - DummyArg, PetscInt, PointerDM, PointerMat, MatReuse, - CallbackPointerIS, CallbackPointerDM, JacobianStruct, - SubMatrixStruct, Initialize, Finalize, ArgvSymbol) +from devito.petsc.types import ( + MultipleFieldData, Initialize, Finalize, ArgvSymbol, MainUserStruct, + CallbackUserStruct +) from devito.petsc.types.macros import petsc_func_begin_user -from devito.petsc.iet.nodes import PetscMetaData -from devito.petsc.utils import core_metadata, petsc_languages -from devito.petsc.iet.routines import (CBBuilder, CCBBuilder, BaseObjectBuilder, - CoupledObjectBuilder, BaseSetup, CoupledSetup, - Solver, CoupledSolver, TimeDependent, - NonTimeDependent) +from devito.petsc.iet.nodes import PetscMetaData, petsc_call +from devito.petsc.config import core_metadata, petsc_languages +from devito.petsc.iet.callbacks import ( + BaseCallbackBuilder, CoupledCallbackBuilder, populate_matrix_context, + get_user_struct_fields +) +from devito.petsc.iet.type_builder import BaseTypeBuilder, CoupledTypeBuilder, objs +from devito.petsc.iet.builder import BuilderBase, CoupledBuilder, make_core_petsc_calls +from devito.petsc.iet.solve import Solve, CoupledSolve +from devito.petsc.iet.time_dependence import TimeDependent, TimeIndependent from devito.petsc.iet.logging import PetscLogger -from devito.petsc.iet.utils import petsc_call, petsc_call_mpi @iet_pass def lower_petsc(iet, **kwargs): - # Check if PETScSolve was used + # Check if `petscsolve` was used inject_solve_mapper = MapNodes(Iteration, PetscMetaData, 'groupby').visit(iet) @@ -52,21 +53,24 @@ def lower_petsc(iet, **kwargs): return finalize(iet), core_metadata() unique_grids = {i.expr.rhs.grid for (i,) in inject_solve_mapper.values()} - # Assumption is that all solves are on the same grid + # Assumption is that all solves are on the same `Grid` if len(unique_grids) > 1: - raise ValueError("All PETScSolves must use the same Grid, but multiple found.") + raise ValueError( + "All `petscsolve` calls must use the same `Grid`, " + "but multiple `Grid`s were found." + ) grid = unique_grids.pop() devito_mpi = kwargs['options'].get('mpi', False) comm = grid.distributor._obj_comm if devito_mpi else 'PETSC_COMM_WORLD' - # Create core PETSc calls (not specific to each PETScSolve) + # Create core PETSc calls (not specific to each `petscsolve`) core = make_core_petsc_calls(objs, comm) setup = [] subs = {} efuncs = {} - # Map PETScSolve to its Section (for logging) + # Map each `PetscMetaData`` to its Section (for logging) section_mapper = MapNodes(Section, PetscMetaData, 'groupby').visit(iet) # Prefixes within the same `Operator` should not be duplicated @@ -76,8 +80,8 @@ def lower_petsc(iet, **kwargs): if duplicates: dup_list = ", ".join(repr(p) for p in sorted(duplicates)) raise ValueError( - f"The following `options_prefix` values are duplicated " - f"among your PETScSolves. Ensure each one is unique: {dup_list}" + "The following `options_prefix` values are duplicated " + f"among your `petscsolve` calls. Ensure each one is unique: {dup_list}" ) # List of `Call`s to clear options from the global PETSc options database, @@ -86,17 +90,17 @@ def lower_petsc(iet, **kwargs): for iters, (inject_solve,) in inject_solve_mapper.items(): - builder = Builder(inject_solve, iters, comm, section_mapper, **kwargs) + solver = BuildSolver(inject_solve, iters, comm, section_mapper, **kwargs) - setup.extend(builder.solver_setup.calls) + setup.extend(solver.builder.calls) # Transform the spatial iteration loop with the calls to execute the solver - subs.update({builder.solve.spatial_body: builder.calls}) + subs.update({solver.solve.spatial_body: solver.calls}) - efuncs.update(builder.cbbuilder.efuncs) + efuncs.update(solver.callback_builder.efuncs) clear_options.extend((petsc_call( - builder.cbbuilder._clear_options_efunc.name, [] + solver.callback_builder._clear_options_efunc.name, [] ),)) populate_matrix_context(efuncs) @@ -108,6 +112,98 @@ def lower_petsc(iet, **kwargs): return iet, metadata +def lower_petsc_symbols(iet, **kwargs): + """ + The `place_definitions` and `place_casts` passes may introduce new + symbols, which must be incorporated into + the relevant PETSc structs. To update the structs, this method then + applies two additional passes: `rebuild_child_user_struct` and + `rebuild_parent_user_struct`. + """ + callback_struct_mapper = {} + # Rebuild `CallbackUserStruct` and update iet accordingly + rebuild_child_user_struct(iet, mapper=callback_struct_mapper) + # Rebuild `MainUserStruct` and update iet accordingly + rebuild_parent_user_struct(iet, mapper=callback_struct_mapper) + + +@iet_pass +def rebuild_child_user_struct(iet, mapper, **kwargs): + """ + Rebuild each `CallbackUserStruct` (the child struct) to include any + new fields introduced by the `place_definitions` and `place_casts` passes. + Also, update the iet accordingly (e.g., dereference the new fields). + - `CallbackUserStruct` is used to access information + in PETSc callback functions via `DMGetApplicationContext`. + """ + old_struct = set([ + i for i in FindSymbols().visit(iet) if isinstance(i, CallbackUserStruct) + ]) + + if not old_struct: + return iet, {} + + # There is a unique `CallbackUserStruct` in each callback + assert len(old_struct) == 1 + old_struct = old_struct.pop() + + # Collect any new fields that have been introduced since the struct was + # previously built + new_fields = [ + f for f in get_user_struct_fields(iet) if f not in old_struct.fields + ] + all_fields = old_struct.fields + new_fields + + # Rebuild the struct + new_struct = old_struct._rebuild(fields=all_fields) + mapper[old_struct] = new_struct + + # Replace old struct with the new one + new_body = Uxreplace(mapper).visit(iet.body) + + # Dereference the new fields and insert them as `standalones` at the top of + # the body. This ensures they are defined before any casts/allocs etc introduced + # by the `place_definitions` and `place_casts` passes. + derefs = tuple([Dereference(i, new_struct) for i in new_fields]) + new_body = new_body._rebuild(standalones=new_body.standalones + derefs) + + return iet._rebuild(body=new_body), {} + + +@iet_pass +def rebuild_parent_user_struct(iet, mapper, **kwargs): + """ + Rebuild each `MainUserStruct` (the parent struct) so that it stays in sync + with its corresponding `CallbackUserStruct` (the child struct). Any IET that + references a parent struct is also updated — either the `PopulateUserContext` + callback or the main Kernel, where the parent struct is registered + via `DMSetApplicationContext`. + """ + if not mapper: + return iet, {} + + parent_struct_mapper = { + v.parent: v.parent._rebuild(fields=v.fields) for v in mapper.values() + } + + if not iet.name.startswith("PopulateUserContext"): + new_body = Uxreplace(parent_struct_mapper).visit(iet.body) + return iet._rebuild(body=new_body), {} + + old_struct = [i for i in iet.parameters if isinstance(i, MainUserStruct)] + assert len(old_struct) == 1 + old_struct = old_struct.pop() + + new_struct = parent_struct_mapper[old_struct] + + new_body = [ + DummyExpr(FieldFromPointer(i._C_symbol, new_struct), i._C_symbol) + for i in new_struct.callback_fields + ] + new_body = iet.body._rebuild(body=new_body) + return iet._rebuild(body=new_body, parameters=(new_struct,)), {} + + def initialize(iet): # should be int because the correct type for argc is a C int # and not a int32 @@ -134,12 +230,7 @@ def finalize(iet): return iet._rebuild(body=finalize_body) -def make_core_petsc_calls(objs, comm): - call_mpi = petsc_call_mpi('MPI_Comm_size', [comm, Byref(objs['size'])]) - return call_mpi, BlankLine - - -class Builder: +class BuildSolver: """ This class is designed to support future extensions, enabling different combinations of solver types, preconditioning methods, @@ -165,39 +256,39 @@ def __init__(self, inject_solve, iters, comm, section_mapper, **kwargs): 'section_mapper': self.section_mapper, **self.kwargs } - self.common_kwargs['solver_objs'] = self.object_builder.solver_objs + self.common_kwargs['solver_objs'] = self.type_builder.solver_objs self.common_kwargs['time_dependence'] = self.time_dependence - self.common_kwargs['cbbuilder'] = self.cbbuilder + self.common_kwargs['callback_builder'] = self.callback_builder self.common_kwargs['logger'] = self.logger @cached_property - def object_builder(self): + def type_builder(self): return ( - CoupledObjectBuilder(**self.common_kwargs) + CoupledTypeBuilder(**self.common_kwargs) if self.coupled else - BaseObjectBuilder(**self.common_kwargs) + BaseTypeBuilder(**self.common_kwargs) ) @cached_property def time_dependence(self): mapper = self.inject_solve.expr.rhs.time_mapper - time_class = TimeDependent if mapper else NonTimeDependent + time_class = TimeDependent if mapper else TimeIndependent return time_class(**self.common_kwargs) @cached_property - def cbbuilder(self): - return CCBBuilder(**self.common_kwargs) \ - if self.coupled else CBBuilder(**self.common_kwargs) + def callback_builder(self): + return CoupledCallbackBuilder(**self.common_kwargs) \ + if self.coupled else BaseCallbackBuilder(**self.common_kwargs) @cached_property - def solver_setup(self): - return CoupledSetup(**self.common_kwargs) \ - if self.coupled else BaseSetup(**self.common_kwargs) + def builder(self): + return CoupledBuilder(**self.common_kwargs) \ + if self.coupled else BuilderBase(**self.common_kwargs) @cached_property def solve(self): - return CoupledSolver(**self.common_kwargs) \ - if self.coupled else Solver(**self.common_kwargs) + return CoupledSolve(**self.common_kwargs) \ + if self.coupled else Solve(**self.common_kwargs) @cached_property def logger(self): @@ -209,81 +300,3 @@ def logger(self): @cached_property def calls(self): return List(body=self.solve.calls+self.logger.calls) - - -def populate_matrix_context(efuncs): - if not objs['dummyefunc'] in efuncs.values(): - return - - subdms_expr = DummyExpr( - FieldFromPointer(objs['Subdms']._C_symbol, objs['ljacctx']), - objs['Subdms']._C_symbol - ) - fields_expr = DummyExpr( - FieldFromPointer(objs['Fields']._C_symbol, objs['ljacctx']), - objs['Fields']._C_symbol - ) - body = CallableBody( - List(body=[subdms_expr, fields_expr]), - init=(petsc_func_begin_user,), - retstmt=tuple([Call('PetscFunctionReturn', arguments=[0])]) - ) - name = 'PopulateMatContext' - efuncs[name] = Callable( - name, body, objs['err'], - parameters=[objs['ljacctx'], objs['Subdms'], objs['Fields']] - ) - - -subdms = PointerDM(name='subdms') -fields = PointerIS(name='fields') -submats = PointerMat(name='submats') -rows = PointerIS(name='rows') -cols = PointerIS(name='cols') - - -# A static dict containing shared symbols and objects that are not -# unique to each PETScSolve. -# Many of these objects are used as arguments in callback functions to make -# the C code cleaner and more modular. This is also a step toward leveraging -# Devito's `reuse_efuncs` functionality, allowing reuse of efuncs when -# they are semantically identical. -objs = frozendict({ - 'size': PetscMPIInt(name='size'), - 'err': PetscErrorCode(name='err'), - 'block': CallbackMat('block'), - 'submat_arr': PointerMat(name='submat_arr'), - 'subblockrows': PetscInt('subblockrows'), - 'subblockcols': PetscInt('subblockcols'), - 'rowidx': PetscInt('rowidx'), - 'colidx': PetscInt('colidx'), - 'J': Mat('J'), - 'X': Vec('X'), - 'xloc': CallbackVec('xloc'), - 'Y': Vec('Y'), - 'yloc': CallbackVec('yloc'), - 'F': Vec('F'), - 'floc': CallbackVec('floc'), - 'B': Vec('B'), - 'nfields': PetscInt('nfields'), - 'irow': PointerIS(name='irow'), - 'icol': PointerIS(name='icol'), - 'nsubmats': Scalar('nsubmats', dtype=np.int32), - 'matreuse': MatReuse('scall'), - 'snes': SNES('snes'), - 'rows': rows, - 'cols': cols, - 'Subdms': subdms, - 'LocalSubdms': CallbackPointerDM(name='subdms'), - 'Fields': fields, - 'LocalFields': CallbackPointerIS(name='fields'), - 'Submats': submats, - 'ljacctx': JacobianStruct( - fields=[subdms, fields, submats], modifier=' *' - ), - 'subctx': SubMatrixStruct(fields=[rows, cols]), - 'dummyctx': Symbol('lctx'), - 'dummyptr': DummyArg('dummy'), - 'dummyefunc': Symbol('dummyefunc'), - 'dof': PetscInt('dof'), -}) diff --git a/devito/petsc/iet/routines.py b/devito/petsc/iet/routines.py deleted file mode 100644 index 94952b3842..0000000000 --- a/devito/petsc/iet/routines.py +++ /dev/null @@ -1,1877 +0,0 @@ -from collections import OrderedDict -from functools import cached_property -import math - -from devito.ir.iet import (Call, FindSymbols, List, Uxreplace, CallableBody, - Dereference, DummyExpr, BlankLine, Callable, FindNodes, - retrieve_iteration_tree, filter_iterations, Iteration, - PointerCast) -from devito.symbolics import (Byref, FieldFromPointer, cast, VOID, - FieldFromComposite, IntDiv, Deref, Mod, String, Null) -from devito.symbolics.unevaluation import Mul -from devito.types.basic import AbstractFunction -from devito.types import Temp, Dimension -from devito.tools import filter_ordered - -from devito.petsc.iet.nodes import (PETScCallable, FormFunctionCallback, - MatShellSetOp, PetscMetaData) -from devito.petsc.iet.utils import (petsc_call, petsc_struct, zero_vector, - dereference_funcs, residual_bundle) -from devito.petsc.types import (PETScArray, PetscBundle, DM, Mat, CallbackVec, Vec, - KSP, PC, SNES, PetscInt, StartPtr, PointerIS, PointerDM, - VecScatter, DMCast, JacobianStruct, SubMatrixStruct, - CallbackDM) -from devito.petsc.types.macros import petsc_func_begin_user - - -class CBBuilder: - """ - Build IET routines to generate PETSc callback functions. - """ - def __init__(self, **kwargs): - - self.rcompile = kwargs.get('rcompile', None) - self.sregistry = kwargs.get('sregistry', None) - self.concretize_mapper = kwargs.get('concretize_mapper', {}) - self.time_dependence = kwargs.get('time_dependence') - self.objs = kwargs.get('objs') - self.solver_objs = kwargs.get('solver_objs') - self.inject_solve = kwargs.get('inject_solve') - self.solve_expr = self.inject_solve.expr.rhs - - self._efuncs = OrderedDict() - self._struct_params = [] - - self._set_options_efunc = None - self._clear_options_efunc = None - self._main_matvec_callback = None - self._user_struct_callback = None - self._F_efunc = None - self._b_efunc = None - - self._J_efuncs = [] - self._initial_guesses = [] - - self._make_core() - self._efuncs = self._uxreplace_efuncs() - - @property - def efuncs(self): - return self._efuncs - - @property - def struct_params(self): - return self._struct_params - - @property - def filtered_struct_params(self): - return filter_ordered(self.struct_params) - - @property - def main_matvec_callback(self): - """ - The matrix-vector callback for the full Jacobian. - This is the function set in the main Kernel via: - PetscCall(MatShellSetOperation(J, MATOP_MULT, (void (*)(void))...)); - The callback has the signature `(Mat, Vec, Vec)`. - """ - return self._J_efuncs[0] - - @property - def J_efuncs(self): - """ - List of matrix-vector callbacks. - Each callback has the signature `(Mat, Vec, Vec)`. Typically, this list - contains a single element, but in mixed systems it can include multiple - callbacks, one for each subblock. - """ - return self._J_efuncs - - @property - def initial_guesses(self): - return self._initial_guesses - - @property - def user_struct_callback(self): - return self._user_struct_callback - - @property - def solver_parameters(self): - return self.solve_expr.solver_parameters - - @property - def field_data(self): - return self.solve_expr.field_data - - @property - def formatted_prefix(self): - return self.solve_expr.formatted_prefix - - @property - def arrays(self): - return self.field_data.arrays - - @property - def target(self): - return self.field_data.target - - def _make_core(self): - self._make_options_callback() - self._make_matvec(self.field_data.jacobian) - self._make_formfunc() - self._make_formrhs() - if self.field_data.initial_guess.exprs: - self._make_initial_guess() - self._make_user_struct_callback() - - def _make_petsc_callable(self, prefix, body, parameters=()): - return PETScCallable( - self.sregistry.make_name(prefix=prefix), - body, - retval=self.objs['err'], - parameters=parameters - ) - - def _make_callable_body(self, body, stacks=(), casts=()): - return CallableBody( - List(body=body), - init=(petsc_func_begin_user,), - stacks=stacks, - casts=casts, - retstmt=(Call('PetscFunctionReturn', arguments=[0]),) - ) - - def _make_options_callback(self): - """ - Create two callbacks: one to set PETSc options and one - to clear them. - - Options are only set/cleared if they were not specifed via - command line arguments. - """ - params = self.solver_parameters - prefix = self.formatted_prefix - - set_body, clear_body = [], [] - - for k, v in params.items(): - option = f'-{prefix}{k}' - - # TODO: Revisit use of a global variable here. - # Consider replacing this with a call to `PetscGetArgs`, though - # initial attempts failed, possibly because the argv pointer is - # created in Python?.. - import devito.petsc.initialize - if option in devito.petsc.initialize._petsc_clargs: - # Ensures that the command line args take priority - continue - - option_name = String(option) - # For options without a value e.g `ksp_view`, pass Null - option_value = Null if v is None else String(str(v)) - set_body.append( - petsc_call('PetscOptionsSetValue', [Null, option_name, option_value]) - ) - clear_body.append( - petsc_call('PetscOptionsClearValue', [Null, option_name]) - ) - - set_body = self._make_callable_body(set_body) - clear_body = self._make_callable_body(clear_body) - - set_callback = self._make_petsc_callable('SetPetscOptions', set_body) - clear_callback = self._make_petsc_callable('ClearPetscOptions', clear_body) - - self._set_options_efunc = set_callback - self._efuncs[set_callback.name] = set_callback - self._clear_options_efunc = clear_callback - self._efuncs[clear_callback.name] = clear_callback - - def _make_matvec(self, jacobian, prefix='MatMult'): - # Compile `matvecs` into an IET via recursive compilation - matvecs = jacobian.matvecs - irs, _ = self.rcompile( - matvecs, options={'mpi': False}, sregistry=self.sregistry, - concretize_mapper=self.concretize_mapper - ) - body = self._create_matvec_body( - List(body=irs.uiet.body), jacobian - ) - objs = self.objs - cb = self._make_petsc_callable( - prefix, body, parameters=(objs['J'], objs['X'], objs['Y']) - ) - self._J_efuncs.append(cb) - self._efuncs[cb.name] = cb - - def _create_matvec_body(self, body, jacobian): - linsolve_expr = self.inject_solve.expr.rhs - objs = self.objs - sobjs = self.solver_objs - - dmda = sobjs['callbackdm'] - ctx = objs['dummyctx'] - xlocal = objs['xloc'] - ylocal = objs['yloc'] - y_matvec = self.arrays[jacobian.row_target]['y'] - x_matvec = self.arrays[jacobian.col_target]['x'] - - body = self.time_dependence.uxreplace_time(body) - - fields = self._dummy_fields(body) - - mat_get_dm = petsc_call('MatGetDM', [objs['J'], Byref(dmda)]) - - dm_get_app_context = petsc_call( - 'DMGetApplicationContext', [dmda, Byref(ctx._C_symbol)] - ) - - zero_y_memory = zero_vector(objs['Y']) if jacobian.zero_memory else None - - dm_get_local_xvec = petsc_call( - 'DMGetLocalVector', [dmda, Byref(xlocal)] - ) - - global_to_local_begin = petsc_call( - 'DMGlobalToLocalBegin', [dmda, objs['X'], - insert_vals, xlocal] - ) - - global_to_local_end = petsc_call('DMGlobalToLocalEnd', [ - dmda, objs['X'], insert_vals, xlocal - ]) - - dm_get_local_yvec = petsc_call( - 'DMGetLocalVector', [dmda, Byref(ylocal)] - ) - - zero_ylocal_memory = zero_vector(ylocal) - - vec_get_array_y = petsc_call( - 'VecGetArray', [ylocal, Byref(y_matvec._C_symbol)] - ) - - vec_get_array_x = petsc_call( - 'VecGetArray', [xlocal, Byref(x_matvec._C_symbol)] - ) - - dm_get_local_info = petsc_call( - 'DMDAGetLocalInfo', [dmda, Byref(linsolve_expr.localinfo)] - ) - - vec_restore_array_y = petsc_call( - 'VecRestoreArray', [ylocal, Byref(y_matvec._C_symbol)] - ) - - vec_restore_array_x = petsc_call( - 'VecRestoreArray', [xlocal, Byref(x_matvec._C_symbol)] - ) - - dm_local_to_global_begin = petsc_call('DMLocalToGlobalBegin', [ - dmda, ylocal, add_vals, objs['Y'] - ]) - - dm_local_to_global_end = petsc_call('DMLocalToGlobalEnd', [ - dmda, ylocal, add_vals, objs['Y'] - ]) - - dm_restore_local_xvec = petsc_call( - 'DMRestoreLocalVector', [dmda, Byref(xlocal)] - ) - - dm_restore_local_yvec = petsc_call( - 'DMRestoreLocalVector', [dmda, Byref(ylocal)] - ) - - # 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, - dm_restore_local_xvec, - dm_restore_local_yvec) - ) - - stacks = ( - mat_get_dm, - dm_get_app_context, - zero_y_memory, - dm_get_local_xvec, - global_to_local_begin, - global_to_local_end, - dm_get_local_yvec, - zero_ylocal_memory, - vec_get_array_y, - vec_get_array_x, - dm_get_local_info - ) - - # Dereference function data in struct - derefs = dereference_funcs(ctx, fields) - - body = self._make_callable_body(body, stacks=stacks+derefs) - - # Replace non-function data with pointer to data in struct - subs = {i._C_symbol: FieldFromPointer(i._C_symbol, ctx) for i in fields} - body = Uxreplace(subs).visit(body) - - self._struct_params.extend(fields) - return body - - def _make_formfunc(self): - objs = self.objs - F_exprs = self.field_data.residual.F_exprs - # Compile `F_exprs` into an IET via recursive compilation - irs, _ = self.rcompile( - F_exprs, options={'mpi': False}, sregistry=self.sregistry, - concretize_mapper=self.concretize_mapper - ) - body_formfunc = self._create_formfunc_body( - List(body=irs.uiet.body) - ) - parameters = (objs['snes'], objs['X'], objs['F'], objs['dummyptr']) - cb = self._make_petsc_callable('FormFunction', body_formfunc, parameters) - - self._F_efunc = cb - self._efuncs[cb.name] = cb - - def _create_formfunc_body(self, body): - linsolve_expr = self.inject_solve.expr.rhs - objs = self.objs - sobjs = self.solver_objs - arrays = self.arrays - target = self.target - - dmda = sobjs['callbackdm'] - ctx = objs['dummyctx'] - - body = self.time_dependence.uxreplace_time(body) - - fields = self._dummy_fields(body) - self._struct_params.extend(fields) - - f_formfunc = arrays[target]['f'] - x_formfunc = arrays[target]['x'] - - dm_cast = DummyExpr(dmda, DMCast(objs['dummyptr']), init=True) - - dm_get_app_context = petsc_call( - 'DMGetApplicationContext', [dmda, Byref(ctx._C_symbol)] - ) - - zero_f_memory = zero_vector(objs['F']) - - dm_get_local_xvec = petsc_call( - 'DMGetLocalVector', [dmda, Byref(objs['xloc'])] - ) - - global_to_local_begin = petsc_call( - 'DMGlobalToLocalBegin', [dmda, objs['X'], - insert_vals, objs['xloc']] - ) - - global_to_local_end = petsc_call('DMGlobalToLocalEnd', [ - dmda, objs['X'], insert_vals, objs['xloc'] - ]) - - dm_get_local_yvec = petsc_call( - 'DMGetLocalVector', [dmda, Byref(objs['floc'])] - ) - - vec_get_array_y = petsc_call( - 'VecGetArray', [objs['floc'], Byref(f_formfunc._C_symbol)] - ) - - vec_get_array_x = petsc_call( - 'VecGetArray', [objs['xloc'], Byref(x_formfunc._C_symbol)] - ) - - dm_get_local_info = petsc_call( - 'DMDAGetLocalInfo', [dmda, Byref(linsolve_expr.localinfo)] - ) - - vec_restore_array_y = petsc_call( - 'VecRestoreArray', [objs['floc'], Byref(f_formfunc._C_symbol)] - ) - - vec_restore_array_x = petsc_call( - 'VecRestoreArray', [objs['xloc'], Byref(x_formfunc._C_symbol)] - ) - - dm_local_to_global_begin = petsc_call('DMLocalToGlobalBegin', [ - dmda, objs['floc'], add_vals, objs['F'] - ]) - - dm_local_to_global_end = petsc_call('DMLocalToGlobalEnd', [ - dmda, objs['floc'], add_vals, objs['F'] - ]) - - dm_restore_local_xvec = petsc_call( - 'DMRestoreLocalVector', [dmda, Byref(objs['xloc'])] - ) - - dm_restore_local_yvec = petsc_call( - 'DMRestoreLocalVector', [dmda, Byref(objs['floc'])] - ) - - body = body._rebuild( - body=body.body + - (vec_restore_array_y, - vec_restore_array_x, - dm_local_to_global_begin, - dm_local_to_global_end, - dm_restore_local_xvec, - dm_restore_local_yvec) - ) - - stacks = ( - dm_cast, - dm_get_app_context, - zero_f_memory, - 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 - ) - - # Dereference function data in struct - derefs = dereference_funcs(ctx, fields) - - body = self._make_callable_body(body, stacks=stacks+derefs) - # Replace non-function data with pointer to data in struct - subs = {i._C_symbol: FieldFromPointer(i._C_symbol, ctx) for i in fields} - - return Uxreplace(subs).visit(body) - - def _make_formrhs(self): - b_exprs = self.field_data.residual.b_exprs - sobjs = self.solver_objs - - # Compile `b_exprs` into an IET via recursive compilation - irs, _ = self.rcompile( - b_exprs, options={'mpi': False}, sregistry=self.sregistry, - concretize_mapper=self.concretize_mapper - ) - body = self._create_form_rhs_body( - List(body=irs.uiet.body) - ) - objs = self.objs - cb = self._make_petsc_callable( - 'FormRHS', body, parameters=(sobjs['callbackdm'], objs['B']) - ) - self._b_efunc = cb - self._efuncs[cb.name] = cb - - def _create_form_rhs_body(self, body): - linsolve_expr = self.inject_solve.expr.rhs - objs = self.objs - sobjs = self.solver_objs - target = self.target - - dmda = sobjs['callbackdm'] - ctx = objs['dummyctx'] - - dm_get_local = petsc_call( - 'DMGetLocalVector', [dmda, Byref(sobjs['blocal'])] - ) - - dm_global_to_local_begin = petsc_call( - 'DMGlobalToLocalBegin', [dmda, objs['B'], - insert_vals, sobjs['blocal']] - ) - - dm_global_to_local_end = petsc_call('DMGlobalToLocalEnd', [ - dmda, objs['B'], insert_vals, - sobjs['blocal'] - ]) - - b_arr = self.field_data.arrays[target]['b'] - - vec_get_array = petsc_call( - 'VecGetArray', [sobjs['blocal'], Byref(b_arr._C_symbol)] - ) - - dm_get_local_info = petsc_call( - 'DMDAGetLocalInfo', [dmda, Byref(linsolve_expr.localinfo)] - ) - - body = self.time_dependence.uxreplace_time(body) - - fields = self._dummy_fields(body) - self._struct_params.extend(fields) - - dm_get_app_context = petsc_call( - 'DMGetApplicationContext', [dmda, Byref(ctx._C_symbol)] - ) - - dm_local_to_global_begin = petsc_call('DMLocalToGlobalBegin', [ - dmda, sobjs['blocal'], insert_vals, - objs['B'] - ]) - - dm_local_to_global_end = petsc_call('DMLocalToGlobalEnd', [ - dmda, sobjs['blocal'], insert_vals, - objs['B'] - ]) - - vec_restore_array = petsc_call( - 'VecRestoreArray', [sobjs['blocal'], Byref(b_arr._C_symbol)] - ) - - dm_restore_local_bvec = petsc_call( - 'DMRestoreLocalVector', [dmda, Byref(sobjs['blocal'])] - ) - - body = body._rebuild(body=body.body + ( - dm_local_to_global_begin, dm_local_to_global_end, vec_restore_array, - dm_restore_local_bvec - )) - - stacks = ( - dm_get_local, - dm_global_to_local_begin, - dm_global_to_local_end, - vec_get_array, - dm_get_app_context, - dm_get_local_info - ) - - # Dereference function data in struct - derefs = dereference_funcs(ctx, fields) - - body = self._make_callable_body([body], stacks=stacks+derefs) - - # Replace non-function data with pointer to data in struct - subs = {i._C_symbol: FieldFromPointer(i._C_symbol, ctx) for - i in fields if not isinstance(i.function, AbstractFunction)} - - return Uxreplace(subs).visit(body) - - def _make_initial_guess(self): - exprs = self.field_data.initial_guess.exprs - sobjs = self.solver_objs - objs = self.objs - - # Compile initital guess `eqns` into an IET via recursive compilation - irs, _ = self.rcompile( - exprs, options={'mpi': False}, sregistry=self.sregistry, - concretize_mapper=self.concretize_mapper - ) - body = self._create_initial_guess_body( - List(body=irs.uiet.body) - ) - cb = self._make_petsc_callable( - 'FormInitialGuess', body, parameters=(sobjs['callbackdm'], objs['xloc']) - ) - self._initial_guesses.append(cb) - self._efuncs[cb.name] = cb - - def _create_initial_guess_body(self, body): - linsolve_expr = self.inject_solve.expr.rhs - objs = self.objs - sobjs = self.solver_objs - target = self.target - - dmda = sobjs['callbackdm'] - ctx = objs['dummyctx'] - - x_arr = self.field_data.arrays[target]['x'] - - vec_get_array = petsc_call( - 'VecGetArray', [objs['xloc'], Byref(x_arr._C_symbol)] - ) - - dm_get_local_info = petsc_call( - 'DMDAGetLocalInfo', [dmda, Byref(linsolve_expr.localinfo)] - ) - - body = self.time_dependence.uxreplace_time(body) - - fields = self._dummy_fields(body) - self._struct_params.extend(fields) - - dm_get_app_context = petsc_call( - 'DMGetApplicationContext', [dmda, Byref(ctx._C_symbol)] - ) - - vec_restore_array = petsc_call( - 'VecRestoreArray', [objs['xloc'], Byref(x_arr._C_symbol)] - ) - - body = body._rebuild(body=body.body + (vec_restore_array,)) - - stacks = ( - vec_get_array, - dm_get_app_context, - dm_get_local_info - ) - - # Dereference function data in struct - derefs = dereference_funcs(ctx, fields) - body = self._make_callable_body(body, stacks=stacks+derefs) - - # Replace non-function data with pointer to data in struct - subs = {i._C_symbol: FieldFromPointer(i._C_symbol, ctx) for - i in fields if not isinstance(i.function, AbstractFunction)} - - return Uxreplace(subs).visit(body) - - def _make_user_struct_callback(self): - """ - This is the struct initialised inside the main kernel and - attached to the DM via DMSetApplicationContext. - # TODO: this could be common between all PETScSolves instead? - """ - mainctx = self.solver_objs['userctx'] = petsc_struct( - self.sregistry.make_name(prefix='ctx'), - self.filtered_struct_params, - self.sregistry.make_name(prefix='UserCtx'), - ) - body = [ - DummyExpr(FieldFromPointer(i._C_symbol, mainctx), i._C_symbol) - for i in mainctx.callback_fields - ] - struct_callback_body = self._make_callable_body(body) - cb = Callable( - self.sregistry.make_name(prefix='PopulateUserContext'), - struct_callback_body, self.objs['err'], - parameters=[mainctx] - ) - self._efuncs[cb.name] = cb - self._user_struct_callback = cb - - def _dummy_fields(self, iet): - # Place all context data required by the shell routines into a struct - fields = [f.function for f in FindSymbols('basics').visit(iet)] - fields = [f for f in fields if not isinstance(f.function, (PETScArray, Temp))] - fields = [ - f for f in fields if not (f.is_Dimension and not (f.is_Time or f.is_Modulo)) - ] - return fields - - def _uxreplace_efuncs(self): - sobjs = self.solver_objs - luserctx = petsc_struct( - sobjs['userctx'].name, - self.filtered_struct_params, - sobjs['userctx'].pname, - modifier=' *' - ) - mapper = {} - visitor = Uxreplace({self.objs['dummyctx']: luserctx}) - for k, v in self._efuncs.items(): - mapper.update({k: visitor.visit(v)}) - return mapper - - -class CCBBuilder(CBBuilder): - def __init__(self, **kwargs): - self._submatrices_callback = None - super().__init__(**kwargs) - - @property - def submatrices_callback(self): - return self._submatrices_callback - - @property - def jacobian(self): - return self.inject_solve.expr.rhs.field_data.jacobian - - @property - def main_matvec_callback(self): - """ - This is the matrix-vector callback associated with the whole Jacobian i.e - is set in the main kernel via - `PetscCall(MatShellSetOperation(J,MATOP_MULT,(void (*)(void))MyMatShellMult));` - """ - return self._main_matvec_callback - - def _make_core(self): - for sm in self.field_data.jacobian.nonzero_submatrices: - self._make_matvec(sm, prefix=f'{sm.name}_MatMult') - - self._make_options_callback() - self._make_whole_matvec() - self._make_whole_formfunc() - self._make_user_struct_callback() - self._create_submatrices() - self._efuncs['PopulateMatContext'] = self.objs['dummyefunc'] - - def _make_whole_matvec(self): - objs = self.objs - body = self._whole_matvec_body() - - parameters = (objs['J'], objs['X'], objs['Y']) - cb = self._make_petsc_callable( - 'WholeMatMult', List(body=body), parameters=parameters - ) - self._main_matvec_callback = cb - self._efuncs[cb.name] = cb - - def _whole_matvec_body(self): - objs = self.objs - sobjs = self.solver_objs - - jctx = objs['ljacctx'] - ctx_main = petsc_call('MatShellGetContext', [objs['J'], Byref(jctx)]) - - nonzero_submats = self.jacobian.nonzero_submatrices - - zero_y_memory = zero_vector(objs['Y']) - - calls = () - for sm in nonzero_submats: - name = sm.name - ctx = sobjs[f'{name}ctx'] - X = sobjs[f'{name}X'] - Y = sobjs[f'{name}Y'] - rows = objs['rows'].base - cols = objs['cols'].base - sm_indexed = objs['Submats'].indexed[sm.linear_idx] - - calls += ( - DummyExpr(sobjs[name], FieldFromPointer(sm_indexed, jctx)), - petsc_call('MatShellGetContext', [sobjs[name], Byref(ctx)]), - petsc_call( - 'VecGetSubVector', - [objs['X'], Deref(FieldFromPointer(cols, ctx)), Byref(X)] - ), - petsc_call( - 'VecGetSubVector', - [objs['Y'], Deref(FieldFromPointer(rows, ctx)), Byref(Y)] - ), - petsc_call('MatMult', [sobjs[name], X, Y]), - petsc_call( - 'VecRestoreSubVector', - [objs['X'], Deref(FieldFromPointer(cols, ctx)), Byref(X)] - ), - petsc_call( - 'VecRestoreSubVector', - [objs['Y'], Deref(FieldFromPointer(rows, ctx)), Byref(Y)] - ), - ) - body = (ctx_main, zero_y_memory, BlankLine) + calls - return self._make_callable_body(body) - - def _make_whole_formfunc(self): - objs = self.objs - F_exprs = self.field_data.residual.F_exprs - # Compile `F_exprs` into an IET via recursive compilation - irs, _ = self.rcompile( - F_exprs, options={'mpi': False}, sregistry=self.sregistry, - concretize_mapper=self.concretize_mapper - ) - body = self._whole_formfunc_body(List(body=irs.uiet.body)) - - parameters = (objs['snes'], objs['X'], objs['F'], objs['dummyptr']) - cb = self._make_petsc_callable( - 'WholeFormFunc', body, parameters=parameters - ) - - self._F_efunc = cb - self._efuncs[cb.name] = cb - - def _whole_formfunc_body(self, body): - linsolve_expr = self.inject_solve.expr.rhs - objs = self.objs - sobjs = self.solver_objs - - dmda = sobjs['callbackdm'] - ctx = objs['dummyctx'] - - body = self.time_dependence.uxreplace_time(body) - - fields = self._dummy_fields(body) - self._struct_params.extend(fields) - - # Process body with bundles for residual callback - bundles = sobjs['bundles'] - fbundle = bundles['f'] - xbundle = bundles['x'] - - body = residual_bundle(body, bundles) - - dm_cast = DummyExpr(dmda, DMCast(objs['dummyptr']), init=True) - - dm_get_app_context = petsc_call( - 'DMGetApplicationContext', [dmda, Byref(ctx._C_symbol)] - ) - - zero_f_memory = zero_vector(objs['F']) - - dm_get_local_xvec = petsc_call( - 'DMGetLocalVector', [dmda, Byref(objs['xloc'])] - ) - - global_to_local_begin = petsc_call( - 'DMGlobalToLocalBegin', [dmda, objs['X'], - insert_vals, objs['xloc']] - ) - - global_to_local_end = petsc_call('DMGlobalToLocalEnd', [ - dmda, objs['X'], insert_vals, objs['xloc'] - ]) - - dm_get_local_yvec = petsc_call( - 'DMGetLocalVector', [dmda, Byref(objs['floc'])] - ) - - vec_get_array_f = petsc_call( - 'VecGetArray', [objs['floc'], Byref(fbundle.vector._C_symbol)] - ) - - vec_get_array_x = petsc_call( - 'VecGetArray', [objs['xloc'], Byref(xbundle.vector._C_symbol)] - ) - - dm_get_local_info = petsc_call( - 'DMDAGetLocalInfo', [dmda, Byref(linsolve_expr.localinfo)] - ) - - vec_restore_array_f = petsc_call( - 'VecRestoreArray', [objs['floc'], Byref(fbundle.vector._C_symbol)] - ) - - vec_restore_array_x = petsc_call( - 'VecRestoreArray', [objs['xloc'], Byref(xbundle.vector._C_symbol)] - ) - - dm_local_to_global_begin = petsc_call('DMLocalToGlobalBegin', [ - dmda, objs['floc'], add_vals, objs['F'] - ]) - - dm_local_to_global_end = petsc_call('DMLocalToGlobalEnd', [ - dmda, objs['floc'], add_vals, objs['F'] - ]) - - dm_restore_local_xvec = petsc_call( - 'DMRestoreLocalVector', [dmda, Byref(objs['xloc'])] - ) - - dm_restore_local_yvec = petsc_call( - 'DMRestoreLocalVector', [dmda, Byref(objs['floc'])] - ) - - body = body._rebuild( - body=body.body + - (vec_restore_array_f, - vec_restore_array_x, - dm_local_to_global_begin, - dm_local_to_global_end, - dm_restore_local_xvec, - dm_restore_local_yvec) - ) - - stacks = ( - dm_cast, - dm_get_app_context, - zero_f_memory, - dm_get_local_xvec, - global_to_local_begin, - global_to_local_end, - dm_get_local_yvec, - vec_get_array_f, - vec_get_array_x, - dm_get_local_info - ) - - # Dereference function data in struct - derefs = dereference_funcs(ctx, fields) - - f_soa = PointerCast(fbundle) - x_soa = PointerCast(xbundle) - - formfunc_body = self._make_callable_body( - body, stacks=stacks+derefs, - casts=(f_soa, x_soa), - ) - # Replace non-function data with pointer to data in struct - subs = {i._C_symbol: FieldFromPointer(i._C_symbol, ctx) for i in fields} - - return Uxreplace(subs).visit(formfunc_body) - - def _create_submatrices(self): - body = self._submat_callback_body() - objs = self.objs - params = ( - objs['J'], - objs['nfields'], - objs['irow'], - objs['icol'], - objs['matreuse'], - objs['Submats'], - ) - cb = self._make_petsc_callable( - 'MatCreateSubMatrices', body, parameters=params) - - self._submatrices_callback = cb - self._efuncs[cb.name] = cb - - def _submat_callback_body(self): - objs = self.objs - sobjs = self.solver_objs - - n_submats = DummyExpr( - objs['nsubmats'], Mul(objs['nfields'], objs['nfields']) - ) - - malloc_submats = petsc_call('PetscCalloc1', [objs['nsubmats'], objs['Submats']]) - - mat_get_dm = petsc_call('MatGetDM', [objs['J'], Byref(sobjs['callbackdm'])]) - - dm_get_app = petsc_call( - 'DMGetApplicationContext', [sobjs['callbackdm'], Byref(objs['dummyctx'])] - ) - - get_ctx = petsc_call('MatShellGetContext', [objs['J'], Byref(objs['ljacctx'])]) - - dm_get_info = petsc_call( - 'DMDAGetInfo', [ - sobjs['callbackdm'], Null, Byref(sobjs['M']), Byref(sobjs['N']), - Null, Null, Null, Null, Byref(objs['dof']), Null, Null, Null, Null, Null - ] - ) - subblock_rows = DummyExpr(objs['subblockrows'], Mul(sobjs['M'], sobjs['N'])) - subblock_cols = DummyExpr(objs['subblockcols'], Mul(sobjs['M'], sobjs['N'])) - - ptr = DummyExpr(objs['submat_arr']._C_symbol, Deref(objs['Submats']), init=True) - - mat_create = petsc_call('MatCreate', [sobjs['comm'], Byref(objs['block'])]) - - mat_set_sizes = petsc_call( - 'MatSetSizes', [ - objs['block'], 'PETSC_DECIDE', 'PETSC_DECIDE', - objs['subblockrows'], objs['subblockcols'] - ] - ) - - mat_set_type = petsc_call('MatSetType', [objs['block'], 'MATSHELL']) - - malloc = petsc_call('PetscMalloc1', [1, Byref(objs['subctx'])]) - i = Dimension(name='i') - - row_idx = DummyExpr(objs['rowidx'], IntDiv(i, objs['dof'])) - col_idx = DummyExpr(objs['colidx'], Mod(i, objs['dof'])) - - deref_subdm = Dereference(objs['Subdms'], objs['ljacctx']) - - set_rows = DummyExpr( - FieldFromPointer(objs['rows'].base, objs['subctx']), - Byref(objs['irow'].indexed[objs['rowidx']]) - ) - set_cols = DummyExpr( - FieldFromPointer(objs['cols'].base, objs['subctx']), - Byref(objs['icol'].indexed[objs['colidx']]) - ) - dm_set_ctx = petsc_call( - 'DMSetApplicationContext', [ - objs['Subdms'].indexed[objs['rowidx']], objs['dummyctx'] - ] - ) - matset_dm = petsc_call('MatSetDM', [ - objs['block'], objs['Subdms'].indexed[objs['rowidx']] - ]) - - set_ctx = petsc_call('MatShellSetContext', [objs['block'], objs['subctx']]) - - mat_setup = petsc_call('MatSetUp', [objs['block']]) - - assign_block = DummyExpr(objs['submat_arr'].indexed[i], objs['block']) - - iter_body = ( - mat_create, - mat_set_sizes, - mat_set_type, - malloc, - row_idx, - col_idx, - set_rows, - set_cols, - dm_set_ctx, - matset_dm, - set_ctx, - mat_setup, - assign_block - ) - - upper_bound = objs['nsubmats'] - 1 - iteration = Iteration(List(body=iter_body), i, upper_bound) - - nonzero_submats = self.jacobian.nonzero_submatrices - matvec_lookup = {mv.name.split('_')[0]: mv for mv in self.J_efuncs} - - matmult_op = [ - petsc_call( - 'MatShellSetOperation', - [ - objs['submat_arr'].indexed[sb.linear_idx], - 'MATOP_MULT', - MatShellSetOp(matvec_lookup[sb.name].name, void, void), - ], - ) - for sb in nonzero_submats if sb.name in matvec_lookup - ] - - body = [ - n_submats, - malloc_submats, - mat_get_dm, - dm_get_app, - dm_get_info, - subblock_rows, - subblock_cols, - ptr, - BlankLine, - iteration, - ] + matmult_op - - return self._make_callable_body(tuple(body), stacks=(get_ctx, deref_subdm)) - - -class BaseObjectBuilder: - """ - A base class for constructing objects needed for a PETSc solver. - Designed to be extended by subclasses, which can override the `_extend_build` - method to support specific use cases. - """ - def __init__(self, **kwargs): - self.inject_solve = kwargs.get('inject_solve') - self.objs = kwargs.get('objs') - self.sregistry = kwargs.get('sregistry') - self.comm = kwargs.get('comm') - self.field_data = self.inject_solve.expr.rhs.field_data - self.solver_objs = self._build() - - def _build(self): - """ - # TODO: update docs - Constructs the core dictionary of solver objects and allows - subclasses to extend or modify it via `_extend_build`. - Returns: - dict: A dictionary containing the following objects: - - 'Jac' (Mat): A matrix representing the jacobian. - - 'xglobal' (GlobalVec): The global solution vector. - - 'xlocal' (LocalVec): The local solution vector. - - 'bglobal': (GlobalVec) Global RHS vector `b`, where `F(x) = b`. - - 'blocal': (LocalVec) Local RHS vector `b`, where `F(x) = b`. - - 'ksp': (KSP) Krylov solver object that manages the linear solver. - - 'pc': (PC) Preconditioner object. - - 'snes': (SNES) Nonlinear solver object. - - 'localsize' (PetscInt): The local length of the solution vector. - - 'dmda' (DM): The DMDA object associated with this solve, linked to - the SNES object via `SNESSetDM`. - - 'callbackdm' (CallbackDM): The DM object accessed within callback - functions via `SNESGetDM`. - """ - sreg = self.sregistry - targets = self.field_data.targets - - snes_name = sreg.make_name(prefix='snes') - formatted_prefix = self.inject_solve.expr.rhs.formatted_prefix - - base_dict = { - 'Jac': Mat(sreg.make_name(prefix='J')), - 'xglobal': Vec(sreg.make_name(prefix='xglobal')), - 'xlocal': Vec(sreg.make_name(prefix='xlocal')), - 'bglobal': Vec(sreg.make_name(prefix='bglobal')), - 'blocal': CallbackVec(sreg.make_name(prefix='blocal')), - 'ksp': KSP(sreg.make_name(prefix='ksp')), - 'pc': PC(sreg.make_name(prefix='pc')), - 'snes': SNES(snes_name), - 'localsize': PetscInt(sreg.make_name(prefix='localsize')), - 'dmda': DM(sreg.make_name(prefix='da'), dofs=len(targets)), - 'callbackdm': CallbackDM(sreg.make_name(prefix='dm')), - 'snes_prefix': String(formatted_prefix), - } - - base_dict['comm'] = self.comm - self._target_dependent(base_dict) - return self._extend_build(base_dict) - - def _target_dependent(self, base_dict): - """ - '_ptr' (StartPtr): A pointer to the beginning of the solution array - that will be updated at each time step. - """ - sreg = self.sregistry - target = self.field_data.target - base_dict[f'{target.name}_ptr'] = StartPtr( - sreg.make_name(prefix=f'{target.name}_ptr'), target.dtype - ) - - def _extend_build(self, base_dict): - """ - Subclasses can override this method to extend or modify the - base dictionary of solver objects. - """ - return base_dict - - -class CoupledObjectBuilder(BaseObjectBuilder): - def _extend_build(self, base_dict): - sreg = self.sregistry - objs = self.objs - targets = self.field_data.targets - arrays = self.field_data.arrays - - base_dict['fields'] = PointerIS( - name=sreg.make_name(prefix='fields'), nindices=len(targets) - ) - base_dict['subdms'] = PointerDM( - name=sreg.make_name(prefix='subdms'), nindices=len(targets) - ) - base_dict['nfields'] = PetscInt(sreg.make_name(prefix='nfields')) - - space_dims = len(self.field_data.grid.dimensions) - - dim_labels = ["M", "N", "P"] - base_dict.update({ - dim_labels[i]: PetscInt(dim_labels[i]) for i in range(space_dims) - }) - - submatrices = self.field_data.jacobian.nonzero_submatrices - - base_dict['jacctx'] = JacobianStruct( - name=sreg.make_name(prefix=objs['ljacctx'].name), - fields=objs['ljacctx'].fields, - ) - - for sm in submatrices: - name = sm.name - base_dict[name] = Mat(name=name) - base_dict[f'{name}ctx'] = SubMatrixStruct( - name=f'{name}ctx', - fields=objs['subctx'].fields, - ) - base_dict[f'{name}X'] = CallbackVec(f'{name}X') - base_dict[f'{name}Y'] = CallbackVec(f'{name}Y') - base_dict[f'{name}F'] = CallbackVec(f'{name}F') - - # Bundle objects/metadata required by the coupled residual callback - f_components, x_components = [], [] - bundle_mapper = {} - pname = sreg.make_name(prefix='Field') - - target_indices = {t: i for i, t in enumerate(targets)} - - for t in targets: - f_arr = arrays[t]['f'] - x_arr = arrays[t]['x'] - f_components.append(f_arr) - x_components.append(x_arr) - - fbundle = PetscBundle( - name='f_bundle', components=f_components, pname=pname - ) - xbundle = PetscBundle( - name='x_bundle', components=x_components, pname=pname - ) - - # Build the bundle mapper - for f_arr, x_arr in zip(f_components, x_components): - bundle_mapper[f_arr.base] = fbundle - bundle_mapper[x_arr.base] = xbundle - - base_dict['bundles'] = { - 'f': fbundle, - 'x': xbundle, - 'bundle_mapper': bundle_mapper, - 'target_indices': target_indices - } - - return base_dict - - def _target_dependent(self, base_dict): - sreg = self.sregistry - targets = self.field_data.targets - for t in targets: - name = t.name - base_dict[f'{name}_ptr'] = StartPtr( - sreg.make_name(prefix=f'{name}_ptr'), t.dtype - ) - base_dict[f'xlocal{name}'] = CallbackVec( - sreg.make_name(prefix=f'xlocal{name}'), liveness='eager' - ) - base_dict[f'Fglobal{name}'] = CallbackVec( - sreg.make_name(prefix=f'Fglobal{name}'), liveness='eager' - ) - base_dict[f'Xglobal{name}'] = CallbackVec( - sreg.make_name(prefix=f'Xglobal{name}') - ) - base_dict[f'xglobal{name}'] = Vec( - sreg.make_name(prefix=f'xglobal{name}') - ) - base_dict[f'blocal{name}'] = CallbackVec( - sreg.make_name(prefix=f'blocal{name}'), liveness='eager' - ) - base_dict[f'bglobal{name}'] = Vec( - sreg.make_name(prefix=f'bglobal{name}') - ) - base_dict[f'da{name}'] = DM( - sreg.make_name(prefix=f'da{name}'), liveness='eager' - ) - base_dict[f'scatter{name}'] = VecScatter( - sreg.make_name(prefix=f'scatter{name}') - ) - - -class BaseSetup: - def __init__(self, **kwargs): - self.inject_solve = kwargs.get('inject_solve') - self.objs = kwargs.get('objs') - self.solver_objs = kwargs.get('solver_objs') - self.cbbuilder = kwargs.get('cbbuilder') - self.field_data = self.inject_solve.expr.rhs.field_data - self.formatted_prefix = self.inject_solve.expr.rhs.formatted_prefix - self.calls = self._setup() - - @property - def snes_ctx(self): - """ - The [optional] context for private data for the function evaluation routine. - https://petsc.org/main/manualpages/SNES/SNESSetFunction/ - """ - return VOID(self.solver_objs['dmda'], stars='*') - - def _setup(self): - sobjs = self.solver_objs - dmda = sobjs['dmda'] - - snes_create = petsc_call('SNESCreate', [sobjs['comm'], Byref(sobjs['snes'])]) - - snes_options_prefix = petsc_call( - 'SNESSetOptionsPrefix', [sobjs['snes'], sobjs['snes_prefix']] - ) if self.formatted_prefix else None - - set_options = petsc_call( - self.cbbuilder._set_options_efunc.name, [] - ) - - snes_set_dm = petsc_call('SNESSetDM', [sobjs['snes'], dmda]) - - create_matrix = petsc_call('DMCreateMatrix', [dmda, Byref(sobjs['Jac'])]) - - snes_set_jac = petsc_call( - 'SNESSetJacobian', [sobjs['snes'], sobjs['Jac'], - sobjs['Jac'], 'MatMFFDComputeJacobian', Null] - ) - - global_x = petsc_call('DMCreateGlobalVector', - [dmda, Byref(sobjs['xglobal'])]) - - target = self.field_data.target - field_from_ptr = FieldFromPointer( - target.function._C_field_data, target.function._C_symbol - ) - - local_size = math.prod( - v for v, dim in zip(target.shape_allocated, target.dimensions) if dim.is_Space - ) - # TODO: Check - VecCreateSeqWithArray - local_x = petsc_call('VecCreateMPIWithArray', - [sobjs['comm'], 1, local_size, 'PETSC_DECIDE', - field_from_ptr, Byref(sobjs['xlocal'])]) - - # TODO: potentially also need to set the DM and local/global map to xlocal - - get_local_size = petsc_call('VecGetSize', - [sobjs['xlocal'], Byref(sobjs['localsize'])]) - - global_b = petsc_call('DMCreateGlobalVector', - [dmda, Byref(sobjs['bglobal'])]) - - snes_get_ksp = petsc_call('SNESGetKSP', - [sobjs['snes'], Byref(sobjs['ksp'])]) - - matvec = self.cbbuilder.main_matvec_callback - matvec_operation = petsc_call( - 'MatShellSetOperation', - [sobjs['Jac'], 'MATOP_MULT', MatShellSetOp(matvec.name, void, void)] - ) - formfunc = self.cbbuilder._F_efunc - formfunc_operation = petsc_call( - 'SNESSetFunction', - [sobjs['snes'], Null, FormFunctionCallback(formfunc.name, void, void), - self.snes_ctx] - ) - - snes_set_options = petsc_call( - 'SNESSetFromOptions', [sobjs['snes']] - ) - - dmda_calls = self._create_dmda_calls(dmda) - - mainctx = sobjs['userctx'] - - call_struct_callback = petsc_call( - self.cbbuilder.user_struct_callback.name, [Byref(mainctx)] - ) - - # TODO: maybe don't need to explictly set this - mat_set_dm = petsc_call('MatSetDM', [sobjs['Jac'], dmda]) - - calls_set_app_ctx = petsc_call('DMSetApplicationContext', [dmda, Byref(mainctx)]) - - base_setup = dmda_calls + ( - snes_create, - snes_options_prefix, - set_options, - snes_set_dm, - create_matrix, - snes_set_jac, - global_x, - local_x, - get_local_size, - global_b, - snes_get_ksp, - matvec_operation, - formfunc_operation, - snes_set_options, - call_struct_callback, - mat_set_dm, - calls_set_app_ctx, - BlankLine - ) - extended_setup = self._extend_setup() - return base_setup + extended_setup - - def _extend_setup(self): - """ - Hook for subclasses to add additional setup calls. - """ - return () - - def _create_dmda_calls(self, dmda): - dmda_create = self._create_dmda(dmda) - dm_setup = petsc_call('DMSetUp', [dmda]) - dm_mat_type = petsc_call('DMSetMatType', [dmda, 'MATSHELL']) - return dmda_create, dm_setup, dm_mat_type - - def _create_dmda(self, dmda): - sobjs = self.solver_objs - grid = self.field_data.grid - nspace_dims = len(grid.dimensions) - - # MPI communicator - args = [sobjs['comm']] - - # Type of ghost nodes - args.extend(['DM_BOUNDARY_GHOSTED' for _ in range(nspace_dims)]) - - # Stencil type - if nspace_dims > 1: - args.append('DMDA_STENCIL_BOX') - - # Global dimensions - args.extend(list(grid.shape)[::-1]) - # No.of processors in each dimension - if nspace_dims > 1: - args.extend(list(grid.distributor.topology)[::-1]) - - # Number of degrees of freedom per node - args.append(dmda.dofs) - # "Stencil width" -> size of overlap - # TODO: Instead, this probably should be - # extracted from field_data.target._size_outhalo? - stencil_width = self.field_data.space_order - - args.append(stencil_width) - args.extend([Null]*nspace_dims) - - # The distributed array object - args.append(Byref(dmda)) - - # The PETSc call used to create the DMDA - dmda = petsc_call(f'DMDACreate{nspace_dims}d', args) - - return dmda - - -class CoupledSetup(BaseSetup): - def _setup(self): - # TODO: minimise code duplication with superclass - objs = self.objs - sobjs = self.solver_objs - dmda = sobjs['dmda'] - - snes_create = petsc_call('SNESCreate', [sobjs['comm'], Byref(sobjs['snes'])]) - - snes_options_prefix = petsc_call( - 'SNESSetOptionsPrefix', [sobjs['snes'], sobjs['snes_prefix']] - ) if self.formatted_prefix else None - - set_options = petsc_call( - self.cbbuilder._set_options_efunc.name, [] - ) - - snes_set_dm = petsc_call('SNESSetDM', [sobjs['snes'], dmda]) - - create_matrix = petsc_call('DMCreateMatrix', [dmda, Byref(sobjs['Jac'])]) - - snes_set_jac = petsc_call( - 'SNESSetJacobian', [sobjs['snes'], sobjs['Jac'], - sobjs['Jac'], 'MatMFFDComputeJacobian', Null] - ) - - global_x = petsc_call('DMCreateGlobalVector', - [dmda, Byref(sobjs['xglobal'])]) - - local_x = petsc_call('DMCreateLocalVector', [dmda, Byref(sobjs['xlocal'])]) - - get_local_size = petsc_call('VecGetSize', - [sobjs['xlocal'], Byref(sobjs['localsize'])]) - - snes_get_ksp = petsc_call('SNESGetKSP', - [sobjs['snes'], Byref(sobjs['ksp'])]) - - matvec = self.cbbuilder.main_matvec_callback - matvec_operation = petsc_call( - 'MatShellSetOperation', - [sobjs['Jac'], 'MATOP_MULT', MatShellSetOp(matvec.name, void, void)] - ) - formfunc = self.cbbuilder._F_efunc - formfunc_operation = petsc_call( - 'SNESSetFunction', - [sobjs['snes'], Null, FormFunctionCallback(formfunc.name, void, void), - self.snes_ctx] - ) - - snes_set_options = petsc_call( - 'SNESSetFromOptions', [sobjs['snes']] - ) - - dmda_calls = self._create_dmda_calls(dmda) - - mainctx = sobjs['userctx'] - - call_struct_callback = petsc_call( - self.cbbuilder.user_struct_callback.name, [Byref(mainctx)] - ) - - # TODO: maybe don't need to explictly set this - mat_set_dm = petsc_call('MatSetDM', [sobjs['Jac'], dmda]) - - calls_set_app_ctx = petsc_call('DMSetApplicationContext', [dmda, Byref(mainctx)]) - - create_field_decomp = petsc_call( - 'DMCreateFieldDecomposition', - [dmda, Byref(sobjs['nfields']), Null, Byref(sobjs['fields']), - Byref(sobjs['subdms'])] - ) - submat_cb = self.cbbuilder.submatrices_callback - matop_create_submats_op = petsc_call( - 'MatShellSetOperation', - [sobjs['Jac'], 'MATOP_CREATE_SUBMATRICES', - MatShellSetOp(submat_cb.name, void, void)] - ) - - call_coupled_struct_callback = petsc_call( - 'PopulateMatContext', - [Byref(sobjs['jacctx']), sobjs['subdms'], sobjs['fields']] - ) - - shell_set_ctx = petsc_call( - 'MatShellSetContext', [sobjs['Jac'], Byref(sobjs['jacctx']._C_symbol)] - ) - - create_submats = petsc_call( - 'MatCreateSubMatrices', - [sobjs['Jac'], sobjs['nfields'], sobjs['fields'], - sobjs['fields'], 'MAT_INITIAL_MATRIX', - Byref(FieldFromComposite(objs['Submats'].base, sobjs['jacctx']))] - ) - - targets = self.field_data.targets - - deref_dms = [ - DummyExpr(sobjs[f'da{t.name}'], sobjs['subdms'].indexed[i]) - for i, t in enumerate(targets) - ] - - xglobals = [petsc_call( - 'DMCreateGlobalVector', - [sobjs[f'da{t.name}'], Byref(sobjs[f'xglobal{t.name}'])] - ) for t in targets] - - coupled_setup = dmda_calls + ( - snes_create, - snes_options_prefix, - set_options, - snes_set_dm, - create_matrix, - snes_set_jac, - global_x, - local_x, - get_local_size, - snes_get_ksp, - matvec_operation, - formfunc_operation, - snes_set_options, - call_struct_callback, - mat_set_dm, - calls_set_app_ctx, - create_field_decomp, - matop_create_submats_op, - call_coupled_struct_callback, - shell_set_ctx, - create_submats) + \ - tuple(deref_dms) + tuple(xglobals) - return coupled_setup - - -class Solver: - def __init__(self, **kwargs): - self.inject_solve = kwargs.get('inject_solve') - self.objs = kwargs.get('objs') - self.solver_objs = kwargs.get('solver_objs') - self.iters = kwargs.get('iters') - self.cbbuilder = kwargs.get('cbbuilder') - self.time_dependence = kwargs.get('time_dependence') - self.calls = self._execute_solve() - - def _execute_solve(self): - """ - Assigns the required time iterators to the struct and executes - the necessary calls to execute the SNES solver. - """ - sobjs = self.solver_objs - target = self.inject_solve.expr.rhs.field_data.target - - struct_assignment = self.time_dependence.assign_time_iters(sobjs['userctx']) - - b_efunc = self.cbbuilder._b_efunc - - dmda = sobjs['dmda'] - - rhs_call = petsc_call(b_efunc.name, [sobjs['dmda'], sobjs['bglobal']]) - - vec_place_array = self.time_dependence.place_array(target) - - if self.cbbuilder.initial_guesses: - initguess = self.cbbuilder.initial_guesses[0] - initguess_call = petsc_call(initguess.name, [dmda, sobjs['xlocal']]) - else: - initguess_call = None - - dm_local_to_global_x = petsc_call( - 'DMLocalToGlobal', [dmda, sobjs['xlocal'], insert_vals, - sobjs['xglobal']] - ) - - snes_solve = petsc_call('SNESSolve', [ - sobjs['snes'], sobjs['bglobal'], sobjs['xglobal']] - ) - - dm_global_to_local_x = petsc_call('DMGlobalToLocal', [ - dmda, sobjs['xglobal'], insert_vals, sobjs['xlocal']] - ) - - vec_reset_array = self.time_dependence.reset_array(target) - - run_solver_calls = (struct_assignment,) + ( - rhs_call, - ) + vec_place_array + ( - initguess_call, - dm_local_to_global_x, - snes_solve, - dm_global_to_local_x, - vec_reset_array, - BlankLine, - ) - return run_solver_calls - - @cached_property - def spatial_body(self): - spatial_body = [] - # TODO: remove the iters[0] - for tree in retrieve_iteration_tree(self.iters[0]): - root = filter_iterations(tree, key=lambda i: i.dim.is_Space) - if root: - root = root[0] - if self.inject_solve in FindNodes(PetscMetaData).visit(root): - spatial_body.append(root) - spatial_body, = spatial_body - return spatial_body - - -class CoupledSolver(Solver): - def _execute_solve(self): - """ - Assigns the required time iterators to the struct and executes - the necessary calls to execute the SNES solver. - """ - sobjs = self.solver_objs - xglob = sobjs['xglobal'] - - struct_assignment = self.time_dependence.assign_time_iters(sobjs['userctx']) - targets = self.inject_solve.expr.rhs.field_data.targets - - # TODO: optimise the ccode generated here - pre_solve = () - post_solve = () - - for i, t in enumerate(targets): - name = t.name - dm = sobjs[f'da{name}'] - target_xloc = sobjs[f'xlocal{name}'] - target_xglob = sobjs[f'xglobal{name}'] - field = sobjs['fields'].indexed[i] - s = sobjs[f'scatter{name}'] - - pre_solve += ( - # TODO: Switch to createwitharray and move to setup - petsc_call('DMCreateLocalVector', [dm, Byref(target_xloc)]), - - # TODO: Need to call reset array - self.time_dependence.place_array(t), - petsc_call( - 'DMLocalToGlobal', - [dm, target_xloc, insert_vals, target_xglob] - ), - petsc_call( - 'VecScatterCreate', - [xglob, field, target_xglob, Null, Byref(s)] - ), - petsc_call( - 'VecScatterBegin', - [s, target_xglob, xglob, insert_vals, sreverse] - ), - petsc_call( - 'VecScatterEnd', - [s, target_xglob, xglob, insert_vals, sreverse] - ), - BlankLine, - ) - - post_solve += ( - petsc_call( - 'VecScatterBegin', - [s, xglob, target_xglob, insert_vals, sforward] - ), - petsc_call( - 'VecScatterEnd', - [s, xglob, target_xglob, insert_vals, sforward] - ), - petsc_call( - 'DMGlobalToLocal', - [dm, target_xglob, insert_vals, target_xloc] - ) - ) - - snes_solve = (petsc_call('SNESSolve', [sobjs['snes'], Null, xglob]),) - - return (struct_assignment,) + pre_solve + snes_solve + post_solve + (BlankLine,) - - -class NonTimeDependent: - def __init__(self, **kwargs): - self.inject_solve = kwargs.get('inject_solve') - self.iters = kwargs.get('iters') - self.sobjs = kwargs.get('solver_objs') - self.kwargs = kwargs - self.origin_to_moddim = self._origin_to_moddim_mapper(self.iters) - self.time_idx_to_symb = self.inject_solve.expr.rhs.time_mapper - - def _origin_to_moddim_mapper(self, iters): - return {} - - def uxreplace_time(self, body): - return body - - def place_array(self, target): - sobjs = self.sobjs - - field_from_ptr = FieldFromPointer( - target.function._C_field_data, target.function._C_symbol - ) - xlocal = sobjs.get(f'xlocal{target.name}', sobjs['xlocal']) - return (petsc_call('VecPlaceArray', [xlocal, field_from_ptr]),) - - def reset_array(self, target): - """ - """ - sobjs = self.sobjs - xlocal = sobjs.get(f'xlocal{target.name}', sobjs['xlocal']) - return ( - petsc_call('VecResetArray', [xlocal]) - ) - - def assign_time_iters(self, struct): - return [] - - -class TimeDependent(NonTimeDependent): - """ - A class for managing time-dependent solvers. - - This includes scenarios where the target is not directly a `TimeFunction`, - but depends on other functions that are. - - Outline of time loop abstraction with PETSc: - - - At PETScSolve, time indices are replaced with temporary `Symbol` objects - via a mapper (e.g., {t: tau0, t + dt: tau1}) to prevent the time loop - from being generated in the callback functions. These callbacks, needed - for each `SNESSolve` at every time step, don't require the time loop, but - may still need access to data from other time steps. - - All `Function` objects are passed through the initial lowering via the - `SolverMetaData` object, ensuring the correct time loop is generated - in the main kernel. - - Another mapper is created based on the modulo dimensions - generated by the `SolverMetaData` object in the main kernel - (e.g., {time: time, t: t0, t + 1: t1}). - - These two mappers are used to generate a final mapper `symb_to_moddim` - (e.g. {tau0: t0, tau1: t1}) which is used at the IET level to - replace the temporary `Symbol` objects in the callback functions with - the correct modulo dimensions. - - Modulo dimensions are updated in the matrix context struct at each time - step and can be accessed in the callback functions where needed. - """ - @property - def time_spacing(self): - return self.inject_solve.expr.rhs.grid.stepping_dim.spacing - - @cached_property - def symb_to_moddim(self): - """ - Maps temporary `Symbol` objects created during `PETScSolve` to their - corresponding modulo dimensions (e.g. creates {tau0: t0, tau1: t1}). - """ - mapper = { - v: k.xreplace({self.time_spacing: 1, -self.time_spacing: -1}) - for k, v in self.time_idx_to_symb.items() - } - return {symb: self.origin_to_moddim[mapper[symb]] for symb in mapper} - - def is_target_time(self, target): - return any(i.is_Time for i in target.dimensions) - - def target_time(self, target): - target_time = [ - i for i, d in zip(target.indices, target.dimensions) - if d.is_Time - ] - assert len(target_time) == 1 - target_time = target_time.pop() - return target_time - - def uxreplace_time(self, body): - return Uxreplace(self.symb_to_moddim).visit(body) - - def _origin_to_moddim_mapper(self, iters): - """ - Creates a mapper of the origin of the time dimensions to their corresponding - modulo dimensions from a list of `Iteration` objects. - - Examples - -------- - >>> iters - (, - ) - >>> _origin_to_moddim_mapper(iters) - {time: time, t: t0, t + 1: t1} - """ - time_iter = [i for i in iters if any(d.is_Time for d in i.dimensions)] - mapper = {} - - if not time_iter: - return mapper - - for i in time_iter: - for d in i.dimensions: - if d.is_Modulo: - mapper[d.origin] = d - elif d.is_Time: - mapper[d] = d - return mapper - - def place_array(self, target): - """ - In the case that the actual target is time-dependent e.g a `TimeFunction`, - a pointer to the first element in the array that will be updated during - the time step is passed to VecPlaceArray(). - - Examples - -------- - >>> target - f1(time + dt, x, y) - >>> calls = place_array(target) - >>> print(List(body=calls)) - float * f1_ptr0 = (time + 1)*localsize0 + (float*)(f1_vec->data); - PetscCall(VecPlaceArray(xlocal0,f1_ptr0)); - - >>> target - f1(t + dt, x, y) - >>> calls = place_array(target) - >>> print(List(body=calls)) - float * f1_ptr0 = t1*localsize0 + (float*)(f1_vec->data); - PetscCall(VecPlaceArray(xlocal0,f1_ptr0)); - """ - sobjs = self.sobjs - - if self.is_target_time(target): - mapper = {self.time_spacing: 1, -self.time_spacing: -1} - - target_time = self.target_time(target).xreplace(mapper) - target_time = self.origin_to_moddim.get(target_time, target_time) - - xlocal = sobjs.get(f'xlocal{target.name}', sobjs['xlocal']) - start_ptr = sobjs[f'{target.name}_ptr'] - - caster = cast(target.dtype, '*') - return ( - DummyExpr( - start_ptr, - caster( - FieldFromPointer(target._C_field_data, target._C_symbol) - ) + Mul(target_time, sobjs['localsize']), - init=True - ), - petsc_call('VecPlaceArray', [xlocal, start_ptr]) - ) - return super().place_array(target) - - def assign_time_iters(self, struct): - """ - Assign required time iterators to the struct. - These iterators are updated at each timestep in the main kernel - for use in callback functions. - - Examples - -------- - >>> struct - ctx - >>> struct.fields - [h_x, x_M, x_m, f1(t, x), t0, t1] - >>> assigned = assign_time_iters(struct) - >>> print(assigned[0]) - ctx.t0 = t0; - >>> print(assigned[1]) - ctx.t1 = t1; - """ - to_assign = [ - f for f in struct.fields if (f.is_Dimension and (f.is_Time or f.is_Modulo)) - ] - time_iter_assignments = [ - DummyExpr(FieldFromComposite(field, struct), field) - for field in to_assign - ] - return time_iter_assignments - - -void = 'void' -insert_vals = 'INSERT_VALUES' -add_vals = 'ADD_VALUES' -sreverse = 'SCATTER_REVERSE' -sforward = 'SCATTER_FORWARD' diff --git a/devito/petsc/iet/solve.py b/devito/petsc/iet/solve.py new file mode 100644 index 0000000000..f6c1fa22d5 --- /dev/null +++ b/devito/petsc/iet/solve.py @@ -0,0 +1,156 @@ +from functools import cached_property + +from devito.ir.iet import ( + BlankLine, FindNodes, retrieve_iteration_tree, filter_iterations +) +from devito.symbolics import Byref, Null + +from devito.petsc.iet.nodes import PetscMetaData, petsc_call +from devito.petsc.types.modes import InsertMode, ScatterMode + + +class Solve: + def __init__(self, **kwargs): + self.inject_solve = kwargs.get('inject_solve') + self.objs = kwargs.get('objs') + self.solver_objs = kwargs.get('solver_objs') + self.iters = kwargs.get('iters') + self.callback_builder = kwargs.get('callback_builder') + self.time_dependence = kwargs.get('time_dependence') + self.calls = self._execute_solve() + + def _execute_solve(self): + """ + Assigns the required time iterators to the struct and executes + the necessary calls to execute the SNES solver. + """ + sobjs = self.solver_objs + target = self.inject_solve.expr.rhs.field_data.target + + struct_assignment = self.time_dependence.assign_time_iters(sobjs['userctx']) + + b_efunc = self.callback_builder._b_efunc + + dmda = sobjs['dmda'] + + rhs_call = petsc_call(b_efunc.name, [sobjs['dmda'], sobjs['bglobal']]) + + vec_place_array = self.time_dependence.place_array(target) + + if self.callback_builder.initial_guesses: + initguess = self.callback_builder.initial_guesses[0] + initguess_call = petsc_call(initguess.name, [dmda, sobjs['xlocal']]) + else: + initguess_call = None + + dm_local_to_global_x = petsc_call( + 'DMLocalToGlobal', [dmda, sobjs['xlocal'], insert_values, + sobjs['xglobal']] + ) + + snes_solve = petsc_call('SNESSolve', [ + sobjs['snes'], sobjs['bglobal'], sobjs['xglobal']] + ) + + dm_global_to_local_x = petsc_call('DMGlobalToLocal', [ + dmda, sobjs['xglobal'], insert_values, sobjs['xlocal']] + ) + + vec_reset_array = self.time_dependence.reset_array(target) + + run_solver_calls = (struct_assignment,) + ( + rhs_call, + ) + vec_place_array + ( + initguess_call, + dm_local_to_global_x, + snes_solve, + dm_global_to_local_x, + vec_reset_array, + BlankLine, + ) + return run_solver_calls + + @cached_property + def spatial_body(self): + spatial_body = [] + # TODO: remove the iters[0] + for tree in retrieve_iteration_tree(self.iters[0]): + root = filter_iterations(tree, key=lambda i: i.dim.is_Space) + if root: + root = root[0] + if self.inject_solve in FindNodes(PetscMetaData).visit(root): + spatial_body.append(root) + spatial_body, = spatial_body + return spatial_body + + +class CoupledSolve(Solve): + def _execute_solve(self): + """ + Assigns the required time iterators to the struct and executes + the necessary calls to execute the SNES solver. + """ + sobjs = self.solver_objs + xglob = sobjs['xglobal'] + + struct_assignment = self.time_dependence.assign_time_iters(sobjs['userctx']) + targets = self.inject_solve.expr.rhs.field_data.targets + + # TODO: optimise the ccode generated here + pre_solve = () + post_solve = () + + for i, t in enumerate(targets): + name = t.name + dm = sobjs[f'da{name}'] + target_xloc = sobjs[f'xlocal{name}'] + target_xglob = sobjs[f'xglobal{name}'] + field = sobjs['fields'].indexed[i] + s = sobjs[f'scatter{name}'] + + pre_solve += ( + # TODO: Need to call reset array + self.time_dependence.place_array(t), + petsc_call( + 'DMLocalToGlobal', + [dm, target_xloc, insert_values, target_xglob] + ), + petsc_call( + 'VecScatterCreate', + [xglob, field, target_xglob, Null, Byref(s)] + ), + petsc_call( + 'VecScatterBegin', + [s, target_xglob, xglob, insert_values, scatter_reverse] + ), + petsc_call( + 'VecScatterEnd', + [s, target_xglob, xglob, insert_values, scatter_reverse] + ), + BlankLine, + ) + + post_solve += ( + petsc_call( + 'VecScatterBegin', + [s, xglob, target_xglob, insert_values, scatter_forward] + ), + petsc_call( + 'VecScatterEnd', + [s, xglob, target_xglob, insert_values, scatter_forward] + ), + petsc_call( + 'DMGlobalToLocal', + [dm, target_xglob, insert_values, target_xloc] + ) + ) + + snes_solve = (petsc_call('SNESSolve', [sobjs['snes'], Null, xglob]),) + + return (struct_assignment,) + pre_solve + snes_solve + post_solve + (BlankLine,) + + +insert_values = InsertMode.insert_values +add_values = InsertMode.add_values +scatter_reverse = ScatterMode.scatter_reverse +scatter_forward = ScatterMode.scatter_forward diff --git a/devito/petsc/iet/time_dependence.py b/devito/petsc/iet/time_dependence.py new file mode 100644 index 0000000000..abc1b7d69e --- /dev/null +++ b/devito/petsc/iet/time_dependence.py @@ -0,0 +1,197 @@ +from functools import cached_property + +from devito.ir.iet import Uxreplace, DummyExpr +from devito.symbolics import FieldFromPointer, cast, FieldFromComposite +from devito.symbolics.unevaluation import Mul + +from devito.petsc.iet.nodes import petsc_call + + +class TimeBase: + def __init__(self, **kwargs): + self.inject_solve = kwargs.get('inject_solve') + self.iters = kwargs.get('iters') + self.sobjs = kwargs.get('solver_objs') + self.kwargs = kwargs + self.origin_to_moddim = self._origin_to_moddim_mapper(self.iters) + self.time_idx_to_symb = self.inject_solve.expr.rhs.time_mapper + + def _origin_to_moddim_mapper(self, iters): + return {} + + def uxreplace_time(self, body): + return body + + def place_array(self, target): + return () + + def reset_array(self, target): + return () + + def assign_time_iters(self, struct): + return [] + + +class TimeIndependent(TimeBase): + pass + + +class TimeDependent(TimeBase): + """ + A class for managing time-dependent solvers. + This includes scenarios where the target is not directly a `TimeFunction`, + but depends on other functions that are. + Outline of time loop abstraction with PETSc: + - At `petscsolve`, time indices are replaced with temporary `Symbol` objects + via a mapper (e.g., {t: tau0, t + dt: tau1}) to prevent the time loop + from being generated in the callback functions. These callbacks, needed + for each `SNESSolve` at every time step, don't require the time loop, but + may still need access to data from other time steps. + - All `Function` objects are passed through the initial lowering via the + `SolverMetaData` object, ensuring the correct time loop is generated + in the main kernel. + - Another mapper is created based on the modulo dimensions + generated by the `SolverMetaData` object in the main kernel + (e.g., {time: time, t: t0, t + 1: t1}). + - These two mappers are used to generate a final mapper `symb_to_moddim` + (e.g. {tau0: t0, tau1: t1}) which is used at the IET level to + replace the temporary `Symbol` objects in the callback functions with + the correct modulo dimensions. + - Modulo dimensions are updated in the matrix context struct at each time + step and can be accessed in the callback functions where needed. + """ + @property + def time_spacing(self): + return self.inject_solve.expr.rhs.grid.stepping_dim.spacing + + @cached_property + def symb_to_moddim(self): + """ + Maps temporary `Symbol` objects created during `petscsolve` to their + corresponding modulo dimensions (e.g. creates {tau0: t0, tau1: t1}). + """ + mapper = { + v: k.xreplace({self.time_spacing: 1, -self.time_spacing: -1}) + for k, v in self.time_idx_to_symb.items() + } + return {symb: self.origin_to_moddim[mapper[symb]] for symb in mapper} + + def is_target_time(self, target): + return any(i.is_Time for i in target.dimensions) + + def target_time(self, target): + target_time = [ + i for i, d in zip(target.indices, target.dimensions) + if d.is_Time + ] + assert len(target_time) == 1 + target_time = target_time.pop() + return target_time + + def uxreplace_time(self, body): + return Uxreplace(self.symb_to_moddim).visit(body) + + def _origin_to_moddim_mapper(self, iters): + """ + Creates a mapper of the origin of the time dimensions to their corresponding + modulo dimensions from a list of `Iteration` objects. + Examples + -------- + >>> iters + (, + ) + >>> _origin_to_moddim_mapper(iters) + {time: time, t: t0, t + 1: t1} + """ + time_iter = [i for i in iters if any(d.is_Time for d in i.dimensions)] + mapper = {} + + if not time_iter: + return mapper + + for i in time_iter: + for d in i.dimensions: + if d.is_Modulo: + mapper[d.origin] = d + elif d.is_Time: + mapper[d] = d + return mapper + + def place_array(self, target): + """ + In the case that the actual target is time-dependent e.g a `TimeFunction`, + a pointer to the first element in the array that will be updated during + the time step is passed to VecPlaceArray(). + Examples + -------- + >>> target + f1(time + dt, x, y) + >>> calls = place_array(target) + >>> print(List(body=calls)) + float * f1_ptr0 = (time + 1)*localsize0 + (float*)(f1_vec->data); + PetscCall(VecPlaceArray(xlocal0,f1_ptr0)); + >>> target + f1(t + dt, x, y) + >>> calls = place_array(target) + >>> print(List(body=calls)) + float * f1_ptr0 = t1*localsize0 + (float*)(f1_vec->data); + PetscCall(VecPlaceArray(xlocal0,f1_ptr0)); + """ + sobjs = self.sobjs + + if self.is_target_time(target): + mapper = {self.time_spacing: 1, -self.time_spacing: -1} + + target_time = self.target_time(target).xreplace(mapper) + target_time = self.origin_to_moddim.get(target_time, target_time) + + xlocal = sobjs.get(f'xlocal{target.name}', sobjs['xlocal']) + start_ptr = sobjs[f'{target.name}_ptr'] + + caster = cast(target.dtype, '*') + return ( + DummyExpr( + start_ptr, + caster( + FieldFromPointer(target._C_field_data, target._C_symbol) + ) + Mul(target_time, sobjs['localsize']), + init=True + ), + petsc_call('VecPlaceArray', [xlocal, start_ptr]) + ) + return super().place_array(target) + + def reset_array(self, target): + if self.is_target_time(target): + sobjs = self.sobjs + xlocal = sobjs.get(f'xlocal{target.name}', sobjs['xlocal']) + return ( + petsc_call('VecResetArray', [xlocal]) + ) + return super().reset_array(target) + + def assign_time_iters(self, struct): + """ + Assign required time iterators to the struct. + These iterators are updated at each timestep in the main kernel + for use in callback functions. + Examples + -------- + >>> struct + ctx + >>> struct.fields + [h_x, x_M, x_m, f1(t, x), t0, t1] + >>> assigned = assign_time_iters(struct) + >>> print(assigned[0]) + ctx.t0 = t0; + >>> print(assigned[1]) + ctx.t1 = t1; + """ + to_assign = [ + f for f in struct.fields if (f.is_Dimension and (f.is_Time or f.is_Modulo)) + ] + time_iter_assignments = [ + DummyExpr(FieldFromComposite(field, struct), field) + for field in to_assign + ] + return time_iter_assignments diff --git a/devito/petsc/iet/type_builder.py b/devito/petsc/iet/type_builder.py new file mode 100644 index 0000000000..8462ebb916 --- /dev/null +++ b/devito/petsc/iet/type_builder.py @@ -0,0 +1,252 @@ +import numpy as np + +from devito.symbolics import String +from devito.types import Symbol +from devito.tools import frozendict + +from devito.petsc.types import ( + PetscBundle, DM, Mat, CallbackVec, Vec, KSP, PC, SNES, PetscInt, StartPtr, + PointerIS, PointerDM, VecScatter, JacobianStruct, SubMatrixStruct, CallbackDM, + PetscMPIInt, PetscErrorCode, PointerMat, MatReuse, CallbackPointerDM, + CallbackPointerIS, CallbackMat, DummyArg, NofSubMats +) + + +class BaseTypeBuilder: + """ + A base class for constructing objects needed for a PETSc solver. + Designed to be extended by subclasses, which can override the `_extend_build` + method to support specific use cases. + """ + def __init__(self, **kwargs): + self.inject_solve = kwargs.get('inject_solve') + self.objs = kwargs.get('objs') + self.sregistry = kwargs.get('sregistry') + self.comm = kwargs.get('comm') + self.field_data = self.inject_solve.expr.rhs.field_data + self.solver_objs = self._build() + + def _build(self): + """ + # TODO: update docs + Constructs the core dictionary of solver objects and allows + subclasses to extend or modify it via `_extend_build`. + Returns: + dict: A dictionary containing the following objects: + - 'Jac' (Mat): A matrix representing the jacobian. + - 'xglobal' (GlobalVec): The global solution vector. + - 'xlocal' (LocalVec): The local solution vector. + - 'bglobal': (GlobalVec) Global RHS vector `b`, where `F(x) = b`. + - 'blocal': (LocalVec) Local RHS vector `b`, where `F(x) = b`. + - 'ksp': (KSP) Krylov solver object that manages the linear solver. + - 'pc': (PC) Preconditioner object. + - 'snes': (SNES) Nonlinear solver object. + - 'localsize' (PetscInt): The local length of the solution vector. + - 'dmda' (DM): The DMDA object associated with this solve, linked to + the SNES object via `SNESSetDM`. + - 'callbackdm' (CallbackDM): The DM object accessed within callback + functions via `SNESGetDM`. + """ + sreg = self.sregistry + targets = self.field_data.targets + + snes_name = sreg.make_name(prefix='snes') + formatted_prefix = self.inject_solve.expr.rhs.formatted_prefix + + base_dict = { + 'Jac': Mat(sreg.make_name(prefix='J')), + 'xglobal': Vec(sreg.make_name(prefix='xglobal')), + 'xlocal': Vec(sreg.make_name(prefix='xlocal')), + 'bglobal': Vec(sreg.make_name(prefix='bglobal')), + 'blocal': CallbackVec(sreg.make_name(prefix='blocal')), + 'ksp': KSP(sreg.make_name(prefix='ksp')), + 'pc': PC(sreg.make_name(prefix='pc')), + 'snes': SNES(snes_name), + 'localsize': PetscInt(sreg.make_name(prefix='localsize')), + 'dmda': DM(sreg.make_name(prefix='da'), dofs=len(targets)), + 'callbackdm': CallbackDM(sreg.make_name(prefix='dm')), + 'snes_prefix': String(formatted_prefix), + } + + base_dict['comm'] = self.comm + self._target_dependent(base_dict) + return self._extend_build(base_dict) + + def _target_dependent(self, base_dict): + """ + '_ptr' (StartPtr): A pointer to the beginning of the solution array + that will be updated at each time step. + """ + sreg = self.sregistry + target = self.field_data.target + base_dict[f'{target.name}_ptr'] = StartPtr( + sreg.make_name(prefix=f'{target.name}_ptr'), target.dtype + ) + + def _extend_build(self, base_dict): + """ + Subclasses can override this method to extend or modify the + base dictionary of solver objects. + """ + return base_dict + + +class CoupledTypeBuilder(BaseTypeBuilder): + def _extend_build(self, base_dict): + sreg = self.sregistry + objs = self.objs + targets = self.field_data.targets + arrays = self.field_data.arrays + + base_dict['fields'] = PointerIS( + name=sreg.make_name(prefix='fields'), nindices=len(targets) + ) + base_dict['subdms'] = PointerDM( + name=sreg.make_name(prefix='subdms'), nindices=len(targets) + ) + base_dict['nfields'] = PetscInt(sreg.make_name(prefix='nfields')) + + space_dims = len(self.field_data.grid.dimensions) + + dim_labels = ["M", "N", "P"] + base_dict.update({ + dim_labels[i]: PetscInt(dim_labels[i]) for i in range(space_dims) + }) + + submatrices = self.field_data.jacobian.nonzero_submatrices + + base_dict['jacctx'] = JacobianStruct( + name=sreg.make_name(prefix=objs['ljacctx'].name), + fields=objs['ljacctx'].fields, + ) + + for sm in submatrices: + name = sm.name + base_dict[name] = Mat(name=name) + base_dict[f'{name}ctx'] = SubMatrixStruct( + name=f'{name}ctx', + fields=objs['subctx'].fields, + ) + base_dict[f'{name}X'] = CallbackVec(f'{name}X') + base_dict[f'{name}Y'] = CallbackVec(f'{name}Y') + base_dict[f'{name}F'] = CallbackVec(f'{name}F') + + # Bundle objects/metadata required by the coupled residual callback + f_components, x_components = [], [] + bundle_mapper = {} + pname = sreg.make_name(prefix='Field') + + target_indices = {t: i for i, t in enumerate(targets)} + + for t in targets: + f_arr = arrays[t]['f'] + x_arr = arrays[t]['x'] + f_components.append(f_arr) + x_components.append(x_arr) + + fbundle = PetscBundle( + name='f_bundle', components=f_components, pname=pname + ) + xbundle = PetscBundle( + name='x_bundle', components=x_components, pname=pname + ) + + # Build the bundle mapper + for f_arr, x_arr in zip(f_components, x_components): + bundle_mapper[f_arr.base] = fbundle + bundle_mapper[x_arr.base] = xbundle + + base_dict['bundles'] = { + 'f': fbundle, + 'x': xbundle, + 'bundle_mapper': bundle_mapper, + 'target_indices': target_indices + } + + return base_dict + + def _target_dependent(self, base_dict): + sreg = self.sregistry + targets = self.field_data.targets + for t in targets: + name = t.name + base_dict[f'{name}_ptr'] = StartPtr( + sreg.make_name(prefix=f'{name}_ptr'), t.dtype + ) + base_dict[f'xlocal{name}'] = CallbackVec( + sreg.make_name(prefix=f'xlocal{name}'), liveness='eager' + ) + base_dict[f'Fglobal{name}'] = CallbackVec( + sreg.make_name(prefix=f'Fglobal{name}'), liveness='eager' + ) + base_dict[f'Xglobal{name}'] = CallbackVec( + sreg.make_name(prefix=f'Xglobal{name}') + ) + base_dict[f'xglobal{name}'] = Vec( + sreg.make_name(prefix=f'xglobal{name}') + ) + base_dict[f'blocal{name}'] = CallbackVec( + sreg.make_name(prefix=f'blocal{name}'), liveness='eager' + ) + base_dict[f'bglobal{name}'] = Vec( + sreg.make_name(prefix=f'bglobal{name}') + ) + base_dict[f'da{name}'] = DM( + sreg.make_name(prefix=f'da{name}'), liveness='eager' + ) + base_dict[f'scatter{name}'] = VecScatter( + sreg.make_name(prefix=f'scatter{name}') + ) + + +subdms = PointerDM(name='subdms') +fields = PointerIS(name='fields') +submats = PointerMat(name='submats') +rows = PointerIS(name='rows') +cols = PointerIS(name='cols') + + +# A static dict containing shared symbols and objects that are not +# unique to each `petscsolve` call. +# Many of these objects are used as arguments in callback functions to make +# the C code cleaner and more modular. +objs = frozendict({ + 'size': PetscMPIInt(name='size'), + 'err': PetscErrorCode(name='err'), + 'block': CallbackMat('block'), + 'submat_arr': PointerMat(name='submat_arr'), + 'subblockrows': PetscInt('subblockrows'), + 'subblockcols': PetscInt('subblockcols'), + 'rowidx': PetscInt('rowidx'), + 'colidx': PetscInt('colidx'), + 'J': Mat('J'), + 'X': Vec('X'), + 'xloc': CallbackVec('xloc'), + 'Y': Vec('Y'), + 'yloc': CallbackVec('yloc'), + 'F': Vec('F'), + 'floc': CallbackVec('floc'), + 'B': Vec('B'), + 'nfields': PetscInt('nfields'), + 'irow': PointerIS(name='irow'), + 'icol': PointerIS(name='icol'), + 'nsubmats': NofSubMats('nsubmats', dtype=np.int32), + # 'nsubmats': PetscInt('nsubmats'), + 'matreuse': MatReuse('scall'), + 'snes': SNES('snes'), + 'rows': rows, + 'cols': cols, + 'Subdms': subdms, + 'LocalSubdms': CallbackPointerDM(name='subdms'), + 'Fields': fields, + 'LocalFields': CallbackPointerIS(name='fields'), + 'Submats': submats, + 'ljacctx': JacobianStruct( + fields=[subdms, fields, submats], modifier=' *' + ), + 'subctx': SubMatrixStruct(fields=[rows, cols]), + 'dummyctx': Symbol('lctx'), + 'dummyptr': DummyArg('dummy'), + 'dummyefunc': Symbol('dummyefunc'), + 'dof': PetscInt('dof'), +}) diff --git a/devito/petsc/iet/utils.py b/devito/petsc/iet/utils.py deleted file mode 100644 index d7ccfb2b4e..0000000000 --- a/devito/petsc/iet/utils.py +++ /dev/null @@ -1,73 +0,0 @@ -from devito.ir.equations import OpPetsc -from devito.ir.iet import Dereference, FindSymbols, Uxreplace -from devito.types.basic import AbstractFunction - -from devito.petsc.iet.nodes import PetscMetaData, PETScCall - - -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, pname, liveness='lazy', modifier=None): - # TODO: Fix this circular import - from devito.petsc.types.object import PETScStruct - return PETScStruct(name=name, pname=pname, - fields=fields, liveness=liveness, - modifier=modifier) - - -def zero_vector(vec): - """ - Set all entries of a PETSc vector to zero. - """ - return petsc_call('VecSet', [vec, 0.0]) - - -def dereference_funcs(struct, fields): - """ - Dereference AbstractFunctions from a struct. - """ - return tuple( - [Dereference(i, struct) for i in - fields if isinstance(i.function, AbstractFunction)] - ) - - -def residual_bundle(body, bundles): - """ - Replaces PetscArrays in `body` with PetscBundle struct field accesses - (e.g., f_v[ix][iy] -> f_bundle[ix][iy].v). - - Example: - f_v[ix][iy] = x_v[ix][iy]; - f_u[ix][iy] = x_u[ix][iy]; - becomes: - f_bundle[ix][iy].v = x_bundle[ix][iy].v; - f_bundle[ix][iy].u = x_bundle[ix][iy].u; - - NOTE: This is used because the data is interleaved for - multi-component DMDAs in PETSc. - """ - mapper = bundles['bundle_mapper'] - indexeds = FindSymbols('indexeds').visit(body) - subs = {} - - for i in indexeds: - if i.base in mapper: - bundle = mapper[i.base] - index = bundles['target_indices'][i.function.target] - index = (index,) + i.indices - subs[i] = bundle.__getitem__(index) - - body = Uxreplace(subs).visit(body) - return 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 = {OpPetsc: PetscMetaData} diff --git a/devito/petsc/initialize.py b/devito/petsc/initialize.py index a4c136f71a..f225caa4b7 100644 --- a/devito/petsc/initialize.py +++ b/devito/petsc/initialize.py @@ -6,6 +6,7 @@ from devito import Operator, switchconfig from devito.types import Symbol from devito.types.equation import PetscEq + from devito.petsc.types import Initialize, Finalize global _petsc_initialized diff --git a/devito/petsc/logging.py b/devito/petsc/logging.py index 979b4b582a..acbf4cc86f 100644 --- a/devito/petsc/logging.py +++ b/devito/petsc/logging.py @@ -7,7 +7,7 @@ from devito.petsc.types import ( PetscInt, PetscScalar, KSPType, KSPConvergedReason, KSPNormType ) -from devito.petsc.utils import petsc_type_to_ctype +from devito.petsc.config import petsc_type_to_ctype class PetscEntry: @@ -29,6 +29,7 @@ def __repr__(self): class PetscSummary(dict): """ + # TODO: Actually print to screen when DEBUG of PERF is enabled A summary of PETSc statistics collected for all solver runs associated with a single operator during execution. """ @@ -49,8 +50,8 @@ def __init__(self, params, *args, **kwargs): # Dynamically create a property on this class for each PETSc function self._add_properties() - # Initialize the summary by adding PETSc information from each PetscInfo - # object (each corresponding to an individual PETScSolve) + # Initialize the summary with PETSc information from each `PetscInfo` + # object (each corresponding to a `petscsolve` call) for i in self.petscinfos: self.add_info(i) @@ -68,8 +69,8 @@ def petsc_entry(self, petscinfo): Create a named tuple entry for the given PetscInfo object, containing the values for each PETSc function call. """ - # Collect the function names associated with this PetscInfo - # instance (i.e., for a single PETScSolve). + # Collect the function names from this `PetscInfo` + # instance (specific to its `petscsolve` call). funcs = [ petsc_return_variable_dict[f].name for f in petscinfo.query_functions ] diff --git a/devito/petsc/solve.py b/devito/petsc/solve.py index 86b6d624e4..3856392436 100644 --- a/devito/petsc/solve.py +++ b/devito/petsc/solve.py @@ -1,19 +1,23 @@ from devito.types.equation import PetscEq -from devito.tools import as_tuple -from devito.petsc.types import (LinearSolverMetaData, PETScArray, DMDALocalInfo, - FieldData, MultipleFieldData, Jacobian, Residual, - MixedResidual, MixedJacobian, InitialGuess) +from devito.tools import filter_ordered, as_tuple +from devito.types import Symbol, SteppingDimension, TimeDimension +from devito.operations.solve import eval_time_derivatives +from devito.symbolics import retrieve_functions, retrieve_dimensions + +from devito.petsc.types import ( + LinearSolverMetaData, PETScArray, DMDALocalInfo, FieldData, MultipleFieldData, + Jacobian, Residual, MixedResidual, MixedJacobian, InitialGuess +) from devito.petsc.types.equation import EssentialBC -from devito.petsc.solver_parameters import (linear_solver_parameters, - format_options_prefix) -from devito.petsc.utils import get_funcs, generate_time_mapper +from devito.petsc.solver_parameters import ( + linear_solver_parameters, format_options_prefix +) -__all__ = ['PETScSolve'] +__all__ = ['petscsolve'] -# TODO: Rename this to petsc_solve, petscsolve? -def PETScSolve(target_exprs, target=None, solver_parameters=None, +def petscsolve(target_exprs, target=None, solver_parameters=None, options_prefix=None, get_info=[]): """ Returns a symbolic expression representing a linear PETSc solver, @@ -31,13 +35,13 @@ def PETScSolve(target_exprs, target=None, solver_parameters=None, - Single-field problem: Pass a single Eq or list of Eq, and specify `target` separately: - PETScSolve(Eq1, target) - PETScSolve([Eq1, Eq2], target) + petscsolve(Eq1, target) + petscsolve([Eq1, Eq2], target) - Multi-field (mixed) problem: Pass a dictionary mapping each target field to its Eq(s): - PETScSolve({u: Eq1, v: Eq2}) - PETScSolve({u: [Eq1, Eq2], v: [Eq3, Eq4]}) + petscsolve({u: Eq1, v: Eq2}) + petscsolve({u: [Eq1, Eq2], v: [Eq3, Eq4]}) target : Function-like The function (e.g., `Function`, `TimeFunction`) into which the linear @@ -120,14 +124,14 @@ def linear_solve_args(self): exprs = as_tuple(exprs) funcs = get_funcs(exprs) - self.time_mapper = generate_time_mapper(funcs) + self.time_mapper = generate_time_mapper(exprs) arrays = self.generate_arrays(target) exprs = sorted(exprs, key=lambda e: not isinstance(e, EssentialBC)) jacobian = Jacobian(target, exprs, arrays, self.time_mapper) residual = Residual(target, exprs, arrays, self.time_mapper, jacobian.scdiag) - initial_guess = InitialGuess(target, exprs, arrays) + initial_guess = InitialGuess(target, exprs, arrays, self.time_mapper) field_data = FieldData( target=target, @@ -137,7 +141,7 @@ def linear_solve_args(self): arrays=arrays ) - return target, tuple(funcs), field_data + return target, funcs, field_data def generate_arrays(self, *targets): return { @@ -162,7 +166,7 @@ def linear_solve_args(self): exprs.extend(e) funcs = get_funcs(exprs) - self.time_mapper = generate_time_mapper(funcs) + self.time_mapper = generate_time_mapper(exprs) targets = list(self.target_exprs.keys()) arrays = self.generate_arrays(*targets) @@ -183,7 +187,47 @@ def linear_solve_args(self): residual=residual ) - return targets[0], tuple(funcs), all_data + return targets[0], funcs, all_data + + +def get_funcs(exprs): + funcs = [ + f for e in exprs + for f in retrieve_functions(eval_time_derivatives(e.lhs - e.rhs)) + ] + return as_tuple(filter_ordered(funcs)) + + +def generate_time_mapper(exprs): + """ + Replace time indices with `Symbols` in expressions used within + PETSc callback functions. These symbols are Uxreplaced at the IET + level to align with the `TimeDimension` and `ModuloDimension` objects + present in the initial lowering. + NOTE: All functions used in PETSc callback functions are attached to + the `SolverMetaData` object, which is passed through the initial lowering + (and subsequently dropped and replaced with calls to run the solver). + Therefore, the appropriate time loop will always be correctly generated inside + the main kernel. + Examples + -------- + >>> exprs = (Eq(f1(t + dt, x, y), g1(t + dt, x, y) + g2(t, x, y)*f1(t, x, y)),) + >>> generate_time_mapper(exprs) + {t + dt: tau0, t: tau1} + """ + # First, map any actual TimeDimensions + time_indices = [d for d in retrieve_dimensions(exprs) if isinstance(d, TimeDimension)] + + funcs = get_funcs(exprs) + + time_indices.extend(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 dict(zip(time_indices, tau_symbs)) localinfo = DMDALocalInfo(name='info', liveness='eager') diff --git a/devito/petsc/solver_parameters.py b/devito/petsc/solver_parameters.py index 7173ec9745..63ea80265b 100644 --- a/devito/petsc/solver_parameters.py +++ b/devito/petsc/solver_parameters.py @@ -1,6 +1,7 @@ -from petsctools import flatten_parameters import itertools +from petsctools import flatten_parameters + # NOTE: Will be extended, the default preconditioner is not going to be 'none' base_solve_defaults = { diff --git a/devito/petsc/types/__init__.py b/devito/petsc/types/__init__.py index f2305a8352..c40b4acf91 100644 --- a/devito/petsc/types/__init__.py +++ b/devito/petsc/types/__init__.py @@ -1,5 +1,6 @@ from .array import * # noqa -from .types import * # noqa +from .metadata import * # noqa from .object import * # noqa from .equation import * # noqa from .macros import * # noqa +from .modes import * # noqa diff --git a/devito/petsc/types/array.py b/devito/petsc/types/array.py index 1bed71ec50..a56490f3af 100644 --- a/devito/petsc/types/array.py +++ b/devito/petsc/types/array.py @@ -20,7 +20,7 @@ class PETScArray(ArrayBasic, Differentiable): PETScArray objects represent vector objects within PETSc. They correspond to the spatial domain of a Function-like object - provided by the user, which is passed to PETScSolve as the target. + provided by the user, which is passed to `petscsolve` as the target. TODO: Potentially re-evaluate and separate into PETScFunction(Differentiable) and then PETScArray(ArrayBasic). diff --git a/devito/petsc/types/equation.py b/devito/petsc/types/equation.py index e819b48a22..fe9611c1fb 100644 --- a/devito/petsc/types/equation.py +++ b/devito/petsc/types/equation.py @@ -6,7 +6,7 @@ class EssentialBC(Eq): """ - Represents an essential boundary condition for use with PETScSolve. + Represents an essential boundary condition for use with `petscsolve`. Due to ongoing work on PetscSection and DMDA integration (WIP), these conditions are imposed as trivial equations. The compiler @@ -16,8 +16,8 @@ class EssentialBC(Eq): Note: - To define an essential boundary condition, use: Eq(target, boundary_value, subdomain=...), - where `target` is the Function-like object passed to PETScSolve. - - SubDomains used for multiple EssentialBCs must not overlap. + where `target` is the Function-like object passed to `petscsolve`. + - SubDomains used for multiple `EssentialBC`s must not overlap. """ pass diff --git a/devito/petsc/types/macros.py b/devito/petsc/types/macros.py index 4355535e64..94d9368b5e 100644 --- a/devito/petsc/types/macros.py +++ b/devito/petsc/types/macros.py @@ -1,5 +1,4 @@ import cgen as c - # TODO: Don't use c.Line here? petsc_func_begin_user = c.Line('PetscFunctionBeginUser;') diff --git a/devito/petsc/types/types.py b/devito/petsc/types/metadata.py similarity index 96% rename from devito/petsc/types/types.py rename to devito/petsc/types/metadata.py index 598fc658fe..d36e088a36 100644 --- a/devito/petsc/types/types.py +++ b/devito/petsc/types/metadata.py @@ -8,7 +8,7 @@ from devito.types.equation import Eq from devito.operations.solve import eval_time_derivatives -from devito.petsc.utils import petsc_variables +from devito.petsc.config import petsc_variables from devito.petsc.types.equation import EssentialBC, ZeroRow, ZeroColumn @@ -212,7 +212,8 @@ def space_dimensions(self): if len(space_dims) > 1: # TODO: This may not actually have to be the case, but enforcing it for now raise ValueError( - "All targets within a PETScSolve have to have the same space dimensions." + "All targets within a `petscsolve` call must have the" + " same space dimensions." ) return space_dims.pop() @@ -222,7 +223,8 @@ def grid(self): grids = [t.grid for t in self.targets] if len(set(grids)) > 1: raise ValueError( - "All targets within a PETScSolve have to have the same grid." + "Multiple `Grid`s detected in `petscsolve`;" + " all targets must share one `Grid`." ) return grids.pop() @@ -236,7 +238,7 @@ def space_order(self): space_orders = [t.space_order for t in self.targets] if len(set(space_orders)) > 1: raise ValueError( - "All targets within a PETScSolve have to have the same space order." + "All targets within a `petscsolve` call must have the same space order." ) return space_orders.pop() @@ -571,9 +573,13 @@ def _make_F_target(self, eq, F_target, targets): # The initial guess satisfies the essential BCs, so this term is zero. # Still included to support Jacobian testing via finite differences. rhs = arrays['x'] - eq.rhs - zero_row = ZeroRow(arrays['f'], rhs, subdomain=eq.subdomain) + zero_row = ZeroRow( + arrays['f'], rhs.subs(self.time_mapper), subdomain=eq.subdomain + ) # Move essential boundary condition to the right-hand side - zero_col = ZeroColumn(arrays['x'], eq.rhs, subdomain=eq.subdomain) + zero_col = ZeroColumn( + arrays['x'], eq.rhs.subs(self.time_mapper), subdomain=eq.subdomain + ) return (zero_row, zero_col) else: @@ -670,9 +676,10 @@ class InitialGuess: symbolic expressions, enforcing the initial guess to satisfy essential boundary conditions. """ - def __init__(self, target, exprs, arrays): + def __init__(self, target, exprs, arrays, time_mapper): self.target = target self.arrays = arrays + self.time_mapper = time_mapper self._build_exprs(as_tuple(exprs)) @property @@ -694,7 +701,7 @@ def _make_initial_guess(self, expr): if isinstance(expr, EssentialBC): assert expr.lhs == self.target return Eq( - self.arrays[self.target]['x'], expr.rhs, + self.arrays[self.target]['x'], expr.rhs.subs(self.time_mapper), subdomain=expr.subdomain ) else: diff --git a/devito/petsc/types/modes.py b/devito/petsc/types/modes.py new file mode 100644 index 0000000000..0850dc38e7 --- /dev/null +++ b/devito/petsc/types/modes.py @@ -0,0 +1,16 @@ +class InsertMode: + """ + How the entries are combined with the current values in the vectors or matrices. + Reference - https://petsc.org/main/manualpages/Sys/InsertMode/ + """ + insert_values = 'INSERT_VALUES' + add_values = 'ADD_VALUES' + + +class ScatterMode: + """ + Determines the direction of a scatter in `VecScatterBegin()` and `VecScatterEnd()`. + Reference - https://petsc.org/release/manualpages/Vec/ScatterMode/ + """ + scatter_reverse = 'SCATTER_REVERSE' + scatter_forward = 'SCATTER_FORWARD' diff --git a/devito/petsc/types/object.py b/devito/petsc/types/object.py index e674a3d014..8db82be365 100644 --- a/devito/petsc/types/object.py +++ b/devito/petsc/types/object.py @@ -1,12 +1,14 @@ from ctypes import POINTER, c_char from devito.tools import CustomDtype, dtype_to_ctype, as_tuple, CustomIntType -from devito.types import (LocalObject, LocalCompositeObject, ModuloDimension, - TimeDimension, ArrayObject, CustomDimension) +from devito.types import ( + LocalObject, LocalCompositeObject, ModuloDimension, TimeDimension, ArrayObject, + CustomDimension, Scalar +) from devito.symbolics import Byref, cast -from devito.types.basic import DataSymbol +from devito.types.basic import DataSymbol, LocalType -from devito.petsc.iet.utils import petsc_call +from devito.petsc.iet.nodes import petsc_call class PetscMixin: @@ -197,7 +199,7 @@ class PETScStruct(LocalCompositeObject): def time_dim_fields(self): """ Fields within the struct that are updated during the time loop. - These are not set in the `PopulateMatContext` callback. + These are not set in the `PopulateUserContext` callback. """ return [f for f in self.fields if isinstance(f, (ModuloDimension, TimeDimension))] @@ -205,7 +207,7 @@ def time_dim_fields(self): @property def callback_fields(self): """ - Fields within the struct that are initialized in the `PopulateMatContext` + Fields within the struct that are initialized in the `PopulateUserContext` callback. These fields are not updated in the time loop. """ return [f for f in self.fields if f not in self.time_dim_fields] @@ -213,6 +215,22 @@ def callback_fields(self): _C_modifier = ' *' +class MainUserStruct(PETScStruct): + pass + + +class CallbackUserStruct(PETScStruct): + __rkwargs__ = PETScStruct.__rkwargs__ + ('parent',) + + def __init__(self, *args, parent=None, **kwargs): + super().__init__(*args, **kwargs) + self._parent = parent + + @property + def parent(self): + return self._parent + + class JacobianStruct(PETScStruct): def __init__(self, name='jctx', pname='JacobianCtx', fields=None, modifier='', liveness='lazy'): @@ -227,7 +245,7 @@ def __init__(self, name='subctx', pname='SubMatrixCtx', fields=None, _C_modifier = None -class PETScArrayObject(PetscMixin, ArrayObject): +class PETScArrayObject(PetscMixin, ArrayObject, LocalType): _data_alignment = False def __init_finalize__(self, *args, **kwargs): @@ -313,6 +331,10 @@ def _C_ctype(self): return POINTER(POINTER(c_char)) +class NofSubMats(Scalar, LocalType): + pass + + FREE_PRIORITY = { PETScArrayObject: 0, Vec: 1, diff --git a/devito/petsc/utils.py b/devito/petsc/utils.py deleted file mode 100644 index a0b5753255..0000000000 --- a/devito/petsc/utils.py +++ /dev/null @@ -1,148 +0,0 @@ -import os -import ctypes -from pathlib import Path - -from devito.tools import memoized_func, filter_ordered -from devito.types import Symbol, SteppingDimension -from devito.operations.solve import eval_time_derivatives -from devito.symbolics import retrieve_functions - - -class PetscOSError(OSError): - pass - - -@memoized_func -def get_petsc_dir(): - petsc_dir = os.environ.get('PETSC_DIR') - if petsc_dir is None: - raise PetscOSError("PETSC_DIR environment variable not set") - else: - petsc_dir = (Path(petsc_dir),) - - petsc_arch = os.environ.get('PETSC_ARCH') - if petsc_arch is not None: - petsc_dir += (petsc_dir[0] / petsc_arch,) - - petsc_installed = petsc_dir[-1] / 'include' / 'petscconf.h' - if not petsc_installed.is_file(): - raise PetscOSError("PETSc is not installed") - - return petsc_dir - - -@memoized_func -def core_metadata(): - petsc_dir = get_petsc_dir() - - petsc_include = tuple([arch / 'include' for arch in petsc_dir]) - petsc_lib = tuple([arch / 'lib' for arch in petsc_dir]) - - return { - 'includes': ('petscsnes.h', 'petscdmda.h'), - 'include_dirs': petsc_include, - 'libs': ('petsc'), - 'lib_dirs': petsc_lib, - 'ldflags': tuple([f"-Wl,-rpath,{lib}" for lib in petsc_lib]) - } - - -@memoized_func -def get_petsc_variables(): - """ - Taken from https://www.firedrakeproject.org/_modules/firedrake/petsc.html - Get a dict of PETSc environment variables from the file: - $PETSC_DIR/$PETSC_ARCH/lib/petsc/conf/petscvariables - """ - try: - petsc_dir = get_petsc_dir() - except PetscOSError: - petsc_variables = {} - else: - path = [petsc_dir[-1], 'lib', 'petsc', 'conf', 'petscvariables'] - variables_path = Path(*path) - - with open(variables_path) as fh: - # Split lines on first '=' (assignment) - splitlines = (line.split("=", maxsplit=1) for line in fh.readlines()) - petsc_variables = {k.strip(): v.strip() for k, v in splitlines} - - return petsc_variables - - -petsc_variables = get_petsc_variables() -# TODO: Use petsctools get_petscvariables() instead? - - -def get_petsc_type_mappings(): - try: - petsc_precision = petsc_variables['PETSC_PRECISION'] - except KeyError: - printer_mapper = {} - petsc_type_to_ctype = {} - else: - petsc_scalar = 'PetscScalar' - # TODO: Check to see whether Petsc is compiled with - # 32-bit or 64-bit integers - printer_mapper = {ctypes.c_int: 'PetscInt'} - - if petsc_precision == 'single': - printer_mapper[ctypes.c_float] = petsc_scalar - elif petsc_precision == 'double': - printer_mapper[ctypes.c_double] = petsc_scalar - - # Used to construct ctypes.Structures that wrap PETSc objects - petsc_type_to_ctype = {v: k for k, v in printer_mapper.items()} - # Add other PETSc types - petsc_type_to_ctype.update({ - 'KSPType': ctypes.c_char_p, - 'KSPConvergedReason': petsc_type_to_ctype['PetscInt'], - 'KSPNormType': petsc_type_to_ctype['PetscInt'], - }) - return printer_mapper, petsc_type_to_ctype - - -petsc_type_mappings, petsc_type_to_ctype = get_petsc_type_mappings() - - -petsc_languages = ['petsc'] - - -def get_funcs(exprs): - funcs = [ - f for e in exprs - for f in retrieve_functions(eval_time_derivatives(e.lhs - e.rhs)) - ] - return filter_ordered(funcs) - - -def generate_time_mapper(funcs): - """ - Replace time indices with `Symbols` in expressions used within - PETSc callback functions. These symbols are Uxreplaced at the IET - level to align with the `TimeDimension` and `ModuloDimension` objects - present in the initial lowering. - NOTE: All functions used in PETSc callback functions are attached to - the `SolverMetaData` object, which is passed through the initial lowering - (and subsequently dropped and replaced with calls to run the solver). - Therefore, the appropriate time loop will always be correctly generated inside - the main kernel. - Examples - -------- - >>> funcs = [ - >>> f1(t + dt, x, y), - >>> g1(t + dt, x, y), - >>> g2(t, x, y), - >>> f1(t, x, y) - >>> ] - >>> generate_time_mapper(funcs) - {t + dt: tau0, t: tau1} - """ - time_indices = list({ - i if isinstance(d, SteppingDimension) else d - 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 dict(zip(time_indices, tau_symbs)) diff --git a/examples/petsc/Poisson/01_poisson.py b/examples/petsc/Poisson/01_poisson.py index 7ed32e8bbd..318e77f000 100644 --- a/examples/petsc/Poisson/01_poisson.py +++ b/examples/petsc/Poisson/01_poisson.py @@ -4,7 +4,7 @@ from devito import (Grid, Function, Eq, Operator, switchconfig, configuration, SubDomain) -from devito.petsc import PETScSolve, EssentialBC +from devito.petsc import petscsolve, EssentialBC from devito.petsc.initialize import PetscInitialize configuration['compiler'] = 'custom' os.environ['CC'] = 'mpicc' @@ -97,7 +97,7 @@ def analytical(x, y): bcs += [EssentialBC(phi, bc, subdomain=sub4)] exprs = [eqn] + bcs - petsc = PETScSolve(exprs, target=phi, solver_parameters={'ksp_rtol': 1e-8}) + petsc = petscsolve(exprs, target=phi, solver_parameters={'ksp_rtol': 1e-8}) with switchconfig(language='petsc'): op = Operator(petsc) diff --git a/examples/petsc/Poisson/02_laplace.py b/examples/petsc/Poisson/02_laplace.py index 9df68f9ab9..7ac621d1db 100644 --- a/examples/petsc/Poisson/02_laplace.py +++ b/examples/petsc/Poisson/02_laplace.py @@ -3,7 +3,7 @@ from devito import (Grid, Function, Eq, Operator, SubDomain, configuration, switchconfig) -from devito.petsc import PETScSolve, EssentialBC +from devito.petsc import petscsolve, EssentialBC from devito.petsc.initialize import PetscInitialize configuration['compiler'] = 'custom' os.environ['CC'] = 'mpicc' @@ -107,7 +107,7 @@ def analytical(x, y, Lx, Ly): bcs += [EssentialBC(phi, bc_func, subdomain=sub4)] # right exprs = [eqn] + bcs - petsc = PETScSolve(exprs, target=phi, solver_parameters={'ksp_rtol': 1e-8}) + petsc = petscsolve(exprs, target=phi, solver_parameters={'ksp_rtol': 1e-8}) with switchconfig(language='petsc'): op = Operator(petsc) diff --git a/examples/petsc/Poisson/03_poisson.py b/examples/petsc/Poisson/03_poisson.py index 9fa9a9e68a..dd72f265ff 100644 --- a/examples/petsc/Poisson/03_poisson.py +++ b/examples/petsc/Poisson/03_poisson.py @@ -4,7 +4,7 @@ from devito import (Grid, Function, Eq, Operator, switchconfig, configuration, SubDomain) -from devito.petsc import PETScSolve, EssentialBC +from devito.petsc import petscsolve, EssentialBC from devito.petsc.initialize import PetscInitialize configuration['compiler'] = 'custom' os.environ['CC'] = 'mpicc' @@ -75,7 +75,7 @@ def analytical(x): bcs += [EssentialBC(u, np.float64(0.), subdomain=sub2)] exprs = [eqn] + bcs - petsc = PETScSolve(exprs, target=u, solver_parameters={'ksp_rtol': 1e-7}) + petsc = petscsolve(exprs, target=u, solver_parameters={'ksp_rtol': 1e-7}) with switchconfig(language='petsc'): op = Operator(petsc) diff --git a/examples/petsc/Poisson/04_poisson.py b/examples/petsc/Poisson/04_poisson.py index 637ce44076..f5f618e5b2 100644 --- a/examples/petsc/Poisson/04_poisson.py +++ b/examples/petsc/Poisson/04_poisson.py @@ -4,7 +4,7 @@ from devito import (Grid, Function, Eq, Operator, switchconfig, configuration, SubDomain) -from devito.petsc import PETScSolve, EssentialBC +from devito.petsc import petscsolve, EssentialBC from devito.petsc.initialize import PetscInitialize configuration['compiler'] = 'custom' os.environ['CC'] = 'mpicc' @@ -104,7 +104,7 @@ def analytical(x, y): bc_eqns += [EssentialBC(u, bcs, subdomain=sub4)] exprs = [eqn]+bc_eqns - petsc = PETScSolve(exprs, target=u, solver_parameters={'ksp_rtol': 1e-6}) + petsc = petscsolve(exprs, target=u, solver_parameters={'ksp_rtol': 1e-6}) with switchconfig(language='petsc'): op = Operator(petsc) diff --git a/examples/petsc/cfd/01_navierstokes.py b/examples/petsc/cfd/01_navierstokes.py index 1c678d977b..13b9e5d450 100644 --- a/examples/petsc/cfd/01_navierstokes.py +++ b/examples/petsc/cfd/01_navierstokes.py @@ -5,7 +5,7 @@ Operator, SubDomain, switchconfig, configuration) from devito.symbolics import retrieve_functions, INT -from devito.petsc import PETScSolve, EssentialBC +from devito.petsc import petscsolve, EssentialBC from devito.petsc.initialize import PetscInitialize configuration['compiler'] = 'custom' os.environ['CC'] = 'mpicc' @@ -248,7 +248,7 @@ def neumann_right(eq, subdomain): bc_pn1 += [neumann_right(neumann_top(eq_pn1, sub8), sub8)] -eqn_p = PETScSolve([eq_pn1]+bc_pn1, pn1.forward) +eqn_p = petscsolve([eq_pn1]+bc_pn1, pn1.forward) eq_u1 = Eq(u1.dt + u1*u1.dxc + v1*u1.dyc, nu*u1.laplace, subdomain=grid.interior) eq_v1 = Eq(v1.dt + u1*v1.dxc + v1*v1.dyc, nu*v1.laplace, subdomain=grid.interior) @@ -285,8 +285,8 @@ def neumann_right(eq, subdomain): bc_petsc_v1 += [EssentialBC(v1.forward, 0., subdomain=sub1)] # top bc_petsc_v1 += [EssentialBC(v1.forward, 0., subdomain=sub2)] # bottom -tentu = PETScSolve([eq_u1]+bc_petsc_u1, u1.forward) -tentv = PETScSolve([eq_v1]+bc_petsc_v1, v1.forward) +tentu = petscsolve([eq_u1]+bc_petsc_u1, u1.forward) +tentv = petscsolve([eq_v1]+bc_petsc_v1, v1.forward) exprs = [tentu, tentv, eqn_p, update_u, update_v] + bc_u1 + bc_v1 diff --git a/examples/petsc/petsc_test.py b/examples/petsc/petsc_test.py index 5d93669d5f..76b0aac957 100644 --- a/examples/petsc/petsc_test.py +++ b/examples/petsc/petsc_test.py @@ -3,7 +3,7 @@ from devito import (Grid, Function, Eq, Operator, configuration, switchconfig) -from devito.petsc import PETScSolve +from devito.petsc import petscsolve from devito.petsc.initialize import PetscInitialize configuration['compiler'] = 'custom' os.environ['CC'] = 'mpicc' @@ -22,7 +22,7 @@ eq = Eq(v, u.laplace, subdomain=grid.interior) -petsc = PETScSolve([eq], u) +petsc = petscsolve([eq], u) with switchconfig(language='petsc'): op = Operator(petsc) diff --git a/examples/petsc/random/01_helmholtz.py b/examples/petsc/random/01_helmholtz.py index 8702fbf298..dc498a0dae 100644 --- a/examples/petsc/random/01_helmholtz.py +++ b/examples/petsc/random/01_helmholtz.py @@ -4,7 +4,7 @@ from devito.symbolics import retrieve_functions, INT from devito import (configuration, Operator, Eq, Grid, Function, SubDomain, switchconfig) -from devito.petsc import PETScSolve +from devito.petsc import petscsolve from devito.petsc.initialize import PetscInitialize configuration['compiler'] = 'custom' os.environ['CC'] = 'mpicc' @@ -241,7 +241,7 @@ def analytical_solution(x, y): bcs += [neumann_left(neumann_top(eqn, sub7), sub7)] bcs += [neumann_right(neumann_top(eqn, sub8), sub8)] - solver = PETScSolve([eqn]+bcs, target=u, solver_parameters={'rtol': 1e-8}) + solver = petscsolve([eqn]+bcs, target=u, solver_parameters={'rtol': 1e-8}) with switchconfig(openmp=False, language='petsc'): op = Operator(solver) diff --git a/examples/petsc/random/02_biharmonic.py b/examples/petsc/random/02_biharmonic.py index f08ffc07de..635c4ff42b 100644 --- a/examples/petsc/random/02_biharmonic.py +++ b/examples/petsc/random/02_biharmonic.py @@ -7,7 +7,7 @@ from devito import (Grid, Function, Eq, Operator, switchconfig, configuration, SubDomain) -from devito.petsc import PETScSolve, EssentialBC +from devito.petsc import petscsolve, EssentialBC from devito.petsc.initialize import PetscInitialize configuration['compiler'] = 'custom' os.environ['CC'] = 'mpicc' @@ -124,7 +124,7 @@ def f_fcn(x, y): # T (see ref) is nonsymmetric so need to set default KSP type to GMRES params = {'ksp_rtol': 1e-10} - petsc = PETScSolve({v: [eqn1]+bc_v, u: [eqn2]+bc_u}, solver_parameters=params) + petsc = petscsolve({v: [eqn1]+bc_v, u: [eqn2]+bc_u}, solver_parameters=params) with switchconfig(language='petsc'): op = Operator(petsc) diff --git a/examples/petsc/seismic/01_staggered_acoustic.py b/examples/petsc/seismic/01_staggered_acoustic.py index fc9e75938d..2352083236 100644 --- a/examples/petsc/seismic/01_staggered_acoustic.py +++ b/examples/petsc/seismic/01_staggered_acoustic.py @@ -2,7 +2,7 @@ import os import numpy as np from examples.seismic.source import DGaussSource, TimeAxis -from devito.petsc import PETScSolve +from devito.petsc import petscsolve from devito.petsc.initialize import PetscInitialize configuration['compiler'] = 'custom' os.environ['CC'] = 'mpicc' @@ -57,12 +57,12 @@ v_x_2 = Eq(vx2.dt, ro * p2.dx) v_z_2 = Eq(vz2.dt, ro * p2.dz) -petsc_v_x_2 = PETScSolve(v_x_2, target=vx2.forward) -petsc_v_z_2 = PETScSolve(v_z_2, target=vz2.forward) +petsc_v_x_2 = petscsolve(v_x_2, target=vx2.forward) +petsc_v_z_2 = petscsolve(v_z_2, target=vz2.forward) p_2 = Eq(p2.dt, l2m * (vx2.forward.dx + vz2.forward.dz)) -petsc_p_2 = PETScSolve(p_2, target=p2.forward, solver_parameters={'ksp_rtol': 1e-7}) +petsc_p_2 = petscsolve(p_2, target=p2.forward, solver_parameters={'ksp_rtol': 1e-7}) with switchconfig(language='petsc'): op_2 = Operator([petsc_v_x_2, petsc_v_z_2, petsc_p_2, src_p_2], opt='noop') @@ -82,12 +82,12 @@ v_x_4 = Eq(vx4.dt, ro * p4.dx) v_z_4 = Eq(vz4.dt, ro * p4.dz) -petsc_v_x_4 = PETScSolve(v_x_4, target=vx4.forward) -petsc_v_z_4 = PETScSolve(v_z_4, target=vz4.forward) +petsc_v_x_4 = petscsolve(v_x_4, target=vx4.forward) +petsc_v_z_4 = petscsolve(v_z_4, target=vz4.forward) p_4 = Eq(p4.dt, l2m * (vx4.forward.dx + vz4.forward.dz)) -petsc_p_4 = PETScSolve(p_4, target=p4.forward, solver_parameters={'ksp_rtol': 1e-7}) +petsc_p_4 = petscsolve(p_4, target=p4.forward, solver_parameters={'ksp_rtol': 1e-7}) with switchconfig(language='petsc'): op_4 = Operator([petsc_v_x_4, petsc_v_z_4, petsc_p_4, src_p_4], opt='noop') diff --git a/tests/test_petsc.py b/tests/test_petsc.py index 28f7c04025..6926ae76cd 100644 --- a/tests/test_petsc.py +++ b/tests/test_petsc.py @@ -3,10 +3,11 @@ import numpy as np import os import re +import sympy as sp from conftest import skipif from devito import (Grid, Function, TimeFunction, Eq, Operator, - configuration, norm, switchconfig, SubDomain) + configuration, norm, switchconfig, SubDomain, sin) from devito.operator.profiling import PerformanceSummary from devito.ir.iet import (Call, ElementalFunction, FindNodes, retrieve_iteration_tree) @@ -16,7 +17,7 @@ PC, KSPConvergedReason, PETScArray, FieldData, MultipleFieldData, SubMatrixBlock) -from devito.petsc.solve import PETScSolve, EssentialBC +from devito.petsc.solve import petscsolve, EssentialBC from devito.petsc.iet.nodes import Expression from devito.petsc.initialize import PetscInitialize from devito.petsc.logging import PetscSummary @@ -129,7 +130,7 @@ def test_petsc_subs(): @skipif('petsc') def test_petsc_solve(): """ - Test PETScSolve. + Test `petscsolve`. """ grid = Grid((2, 2), dtype=np.float64) @@ -138,7 +139,7 @@ def test_petsc_solve(): eqn = Eq(f.laplace, g) - petsc = PETScSolve(eqn, f) + petsc = petscsolve(eqn, f) with switchconfig(language='petsc'): op = Operator(petsc, opt='noop') @@ -178,7 +179,7 @@ def test_petsc_solve(): @skipif('petsc') def test_multiple_petsc_solves(): """ - Test multiple PETScSolves. + Test multiple `petscsolve` calls, passed to a single `Operator`. """ grid = Grid((2, 2), dtype=np.float64) @@ -191,8 +192,8 @@ def test_multiple_petsc_solves(): eqn1 = Eq(f1.laplace, g1) eqn2 = Eq(f2.laplace, g2) - petsc1 = PETScSolve(eqn1, f1, options_prefix='pde1') - petsc2 = PETScSolve(eqn2, f2, options_prefix='pde2') + petsc1 = petscsolve(eqn1, f1, options_prefix='pde1') + petsc2 = petscsolve(eqn2, f2, options_prefix='pde2') with switchconfig(language='petsc'): op = Operator([petsc1, petsc2], opt='noop') @@ -222,9 +223,9 @@ def test_petsc_cast(): eqn2 = Eq(f2.laplace, 10) eqn3 = Eq(f3.laplace, 10) - petsc1 = PETScSolve(eqn1, f1) - petsc2 = PETScSolve(eqn2, f2) - petsc3 = PETScSolve(eqn3, f3) + petsc1 = petscsolve(eqn1, f1) + petsc2 = petscsolve(eqn2, f2) + petsc3 = petscsolve(eqn3, f3) with switchconfig(language='petsc'): op1 = Operator(petsc1) @@ -254,9 +255,9 @@ def test_dmda_create(): eqn2 = Eq(f2.laplace, 10) eqn3 = Eq(f3.laplace, 10) - petsc1 = PETScSolve(eqn1, f1) - petsc2 = PETScSolve(eqn2, f2) - petsc3 = PETScSolve(eqn3, f3) + petsc1 = petscsolve(eqn1, f1) + petsc2 = petscsolve(eqn2, f2) + petsc3 = petscsolve(eqn3, f3) with switchconfig(language='petsc'): op1 = Operator(petsc1, opt='noop') @@ -275,34 +276,113 @@ def test_dmda_create(): ',1,1,1,1,6,NULL,NULL,NULL,&da0));' in str(op3) -@skipif('petsc') -def test_cinterface_petsc_struct(): +class TestStruct: + @skipif('petsc') + def test_cinterface_petsc_struct(self): - grid = Grid(shape=(11, 11), dtype=np.float64) - f = Function(name='f', grid=grid, space_order=2) - eq = Eq(f.laplace, 10) - petsc = PETScSolve(eq, f) + grid = Grid(shape=(11, 11), dtype=np.float64) + f = Function(name='f', grid=grid, space_order=2) + eq = Eq(f.laplace, 10) + petsc = petscsolve(eq, f) - name = "foo" + name = "foo" - with switchconfig(language='petsc'): - op = Operator(petsc, name=name) + with switchconfig(language='petsc'): + 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 UserCtx` only appears in the header file + assert 'struct UserCtx0\n{' not in ccode + assert 'struct UserCtx0\n{' in hcode + + @skipif('petsc') + def test_temp_arrays_in_struct(self): + + grid = Grid(shape=(11, 11, 11), dtype=np.float64) + + u = TimeFunction(name='u', grid=grid, space_order=2) + x, y, _ = grid.dimensions + + eqn = Eq(u.forward, sin(sp.pi*(x+y)/3.), subdomain=grid.interior) + petsc = petscsolve(eqn, target=u.forward) + + with switchconfig(log_level='DEBUG', language='petsc'): + op = Operator(petsc) + # Check that it runs + op.apply(time_M=3) + + assert 'ctx0->x_size = x_size;' in str(op.ccode) + assert 'ctx0->y_size = y_size;' in str(op.ccode) + + assert 'const PetscInt y_size = ctx0->y_size;' in str(op.ccode) + assert 'const PetscInt x_size = ctx0->x_size;' in str(op.ccode) + + @skipif('petsc') + def test_parameters(self): + + grid = Grid((2, 2), dtype=np.float64) + + 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) - # Trigger the generation of a .c and a .h files - ccode, hcode = op.cinterface(force=True) + eqn2 = Eq(f1, g1*mu2) - 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)) + with switchconfig(language='petsc'): + 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, LocalCompositeObject) for i in op.parameters) - ccode = str(ccode) - hcode = str(hcode) + @skipif('petsc') + def test_field_order(self): + """Verify that the order of fields in the user struct is fixed for + `identical` Operator instances. + """ + grid = Grid(shape=(11, 11, 11), dtype=np.float64) + f = TimeFunction(name='f', grid=grid, space_order=2) + x, y, _ = grid.dimensions + t = grid.time_dim + eq = Eq(f.dt, f.laplace + t*0.005 + sin(sp.pi*(x+y)/3.), subdomain=grid.interior) + petsc = petscsolve(eq, f.forward) - assert 'include "%s.h"' % name in ccode + with switchconfig(language='petsc'): + op1 = Operator(petsc, name="foo1") + op2 = Operator(petsc, name="foo2") - # The public `struct UserCtx` only appears in the header file - assert 'struct UserCtx0\n{' not in ccode - assert 'struct UserCtx0\n{' in hcode + op1_user_struct = op1._func_table['PopulateUserContext0'].root.parameters[0] + op2_user_struct = op2._func_table['PopulateUserContext0'].root.parameters[0] + + assert len(op1_user_struct.fields) == len(op2_user_struct.fields) + assert len(op1_user_struct.callback_fields) == \ + len(op1_user_struct.callback_fields) + assert str(op1_user_struct.fields) == str(op2_user_struct.fields) @skipif('petsc') @@ -317,7 +397,7 @@ def test_callback_arguments(): eqn1 = Eq(f1.laplace, g1) - petsc1 = PETScSolve(eqn1, f1) + petsc1 = petscsolve(eqn1, f1) with switchconfig(language='petsc'): op = Operator(petsc1) @@ -332,39 +412,6 @@ def test_callback_arguments(): assert str(ff.parameters) == '(snes, X, F, dummy)' -@skipif('petsc') -def test_petsc_struct(): - - grid = Grid((2, 2), dtype=np.float64) - - 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(language='petsc'): - 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, LocalCompositeObject) for i in op.parameters) - - @skipif('petsc') def test_apply(): @@ -376,7 +423,7 @@ def test_apply(): eqn = Eq(pn.laplace*mu, rhs, subdomain=grid.interior) - petsc = PETScSolve(eqn, pn) + petsc = petscsolve(eqn, pn) with switchconfig(language='petsc'): # Build the op @@ -399,7 +446,7 @@ def test_petsc_frees(): g = Function(name='g', grid=grid, space_order=2) eqn = Eq(f.laplace, g) - petsc = PETScSolve(eqn, f) + petsc = petscsolve(eqn, f) with switchconfig(language='petsc'): op = Operator(petsc) @@ -424,7 +471,7 @@ def test_calls_to_callbacks(): g = Function(name='g', grid=grid, space_order=2) eqn = Eq(f.laplace, g) - petsc = PETScSolve(eqn, f) + petsc = petscsolve(eqn, f) with switchconfig(language='petsc'): op = Operator(petsc) @@ -448,7 +495,7 @@ def test_start_ptr(): grid = Grid((11, 11), dtype=np.float64) u1 = TimeFunction(name='u1', grid=grid, space_order=2) eq1 = Eq(u1.dt, u1.laplace, subdomain=grid.interior) - petsc1 = PETScSolve(eq1, u1.forward) + petsc1 = petscsolve(eq1, u1.forward) with switchconfig(language='petsc'): op1 = Operator(petsc1) @@ -460,7 +507,7 @@ def test_start_ptr(): # Verify the case with no modulo time stepping u2 = TimeFunction(name='u2', grid=grid, space_order=2, save=5) eq2 = Eq(u2.dt, u2.laplace, subdomain=grid.interior) - petsc2 = PETScSolve(eq2, u2.forward) + petsc2 = petscsolve(eq2, u2.forward) with switchconfig(language='petsc'): op2 = Operator(petsc2) @@ -469,86 +516,146 @@ def test_start_ptr(): '(PetscScalar*)(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), dtype=np.float64) +class TestTimeLoop: + @skipif('petsc') + @pytest.mark.parametrize('dim', [1, 2, 3]) + def test_time_dimensions(self, dim): + """ + 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. + """ + shape = tuple(11 for _ in range(dim)) + grid = Grid(shape=shape, dtype=np.float64) - # 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) + # Modulo time stepping + u1 = TimeFunction(name='u1', grid=grid, space_order=2) + v1 = Function(name='v1', grid=grid, space_order=2) + eq1 = Eq(v1.laplace, u1) + petsc1 = petscsolve(eq1, v1) - with switchconfig(language='petsc'): - op1 = Operator(petsc1) - op1.apply(time_M=3) - body1 = str(op1.body) - rhs1 = str(op1._func_table['FormRHS0'].root.ccode) + with switchconfig(language='petsc'): + op1 = Operator(petsc1) + op1.apply(time_M=3) + body1 = str(op1.body) + rhs1 = str(op1._func_table['FormRHS0'].root.ccode) + + assert 'ctx0.t0 = t0' in body1 + assert 'ctx0.t1 = t1' not in body1 + assert 'ctx0->t0' in rhs1 + assert 'ctx0->t1' not in rhs1 + + # Non-modulo time stepping + u2 = TimeFunction(name='u2', grid=grid, space_order=2, save=5) + v2 = Function(name='v2', grid=grid, space_order=2, save=5) + eq2 = Eq(v2.laplace, u2) + petsc2 = petscsolve(eq2, v2) - assert 'ctx0.t0 = t0' in body1 - assert 'ctx0.t1 = t1' not in body1 - assert 'ctx0->t0' in rhs1 - assert 'ctx0->t1' not in rhs1 + with switchconfig(language='petsc'): + op2 = Operator(petsc2) + op2.apply(time_M=3) + body2 = str(op2.body) + rhs2 = str(op2._func_table['FormRHS0'].root.ccode) - # 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) + assert 'ctx0.time = time' in body2 + assert 'ctx0->time' in rhs2 - with switchconfig(language='petsc'): - op2 = Operator(petsc2) - op2.apply(time_M=3) - body2 = str(op2.body) - rhs2 = str(op2._func_table['FormRHS0'].root.ccode) + # 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) - assert 'ctx0.time = time' in body2 - assert 'ctx0->time' in rhs2 + with switchconfig(language='petsc'): + op3 = Operator(petsc3) + op3.apply(time_M=3) + body3 = str(op3.body) + rhs3 = str(op3._func_table['FormRHS0'].root.ccode) + + assert 'ctx0.t0 = t0' in body3 + assert 'ctx0.t1 = t1' in body3 + assert 'ctx0->t0' in rhs3 + assert 'ctx0->t1' in rhs3 + + # Multiple petsc solves within the same time loop + v2 = Function(name='v2', grid=grid, space_order=2) + eq4 = Eq(v1.laplace, u1) + petsc4 = petscsolve(eq4, v1) + eq5 = Eq(v2.laplace, u1) + petsc5 = petscsolve(eq5, v2) - # 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) + with switchconfig(language='petsc'): + op4 = Operator([petsc4, petsc5]) + op4.apply(time_M=3) + body4 = str(op4.body) - with switchconfig(language='petsc'): - op3 = Operator(petsc3) - op3.apply(time_M=3) - body3 = str(op3.body) - rhs3 = str(op3._func_table['FormRHS0'].root.ccode) - - assert 'ctx0.t0 = t0' in body3 - assert 'ctx0.t1 = t1' in body3 - assert 'ctx0->t0' in rhs3 - assert 'ctx0->t1' in rhs3 - - # Multiple petsc solves within the same time loop - v2 = Function(name='v2', grid=grid, space_order=2) - eq4 = Eq(v1.laplace, u1) - petsc4 = PETScSolve(eq4, v1) - eq5 = Eq(v2.laplace, u1) - petsc5 = PETScSolve(eq5, v2) + assert 'ctx0.t0 = t0' in body4 + assert body4.count('ctx0.t0 = t0') == 1 - with switchconfig(language='petsc'): - op4 = Operator([petsc4, petsc5]) - op4.apply(time_M=3) - body4 = str(op4.body) + @skipif('petsc') + @pytest.mark.parametrize('dim', [1, 2, 3]) + def test_trivial_operator(self, dim): + """ + Test trivial time-dependent problems with `petscsolve`. + """ + # create shape based on dimension + shape = tuple(4 for _ in range(dim)) + grid = Grid(shape=shape, dtype=np.float64) + u = TimeFunction(name='u', grid=grid, save=3) + + eqn = Eq(u.forward, u + 1) - assert 'ctx0.t0 = t0' in body4 - assert body4.count('ctx0.t0 = t0') == 1 + petsc = petscsolve(eqn, target=u.forward) + + with switchconfig(log_level='DEBUG'): + op = Operator(petsc, language='petsc') + op.apply() + + assert np.all(u.data[0] == 0.) + assert np.all(u.data[1] == 1.) + assert np.all(u.data[2] == 2.) + + @skipif('petsc') + @pytest.mark.parametrize('dim', [1, 2, 3]) + def test_time_dim(self, dim): + """ + Verify the time loop abstraction + when a mixture of TimeDimensions and time dependent + SteppingDimensions are used + """ + shape = tuple(4 for _ in range(dim)) + grid = Grid(shape=shape, dtype=np.float64) + # Use modoulo time stepping, i.e don't pass the save argument + u = TimeFunction(name='u', grid=grid) + # Use grid.time_dim in the equation, as well as the TimeFunction itself + petsc = petscsolve(Eq(u.forward, u + 1 + grid.time_dim), target=u.forward) + + with switchconfig(): + op = Operator(petsc, language='petsc') + op.apply(time_M=1) + + body = str(op.body) + rhs = str(op._func_table['FormRHS0'].root.ccode) + + # Check both ctx0.t0 and ctx0.time are assigned since they are both used + # in the callback functions, specifically in FormRHS0 + assert 'ctx0.t0 = t0' in body + assert 'ctx0.time = time' in body + assert 'ctx0->t0' in rhs + assert 'ctx0->time' in rhs + + # Check the ouput is as expected given two time steps have been + # executed (time_M=1) + assert np.all(u.data[1] == 1.) + assert np.all(u.data[0] == 3.) @skipif('petsc') def test_solve_output(): """ - Verify that PETScSolve returns the correct output for - simple cases e.g with the identity matrix. + Verify that `petscsolve` returns the correct output for + simple cases e.g. forming the identity matrix. """ grid = Grid(shape=(11, 11), dtype=np.float64) @@ -558,7 +665,7 @@ def test_solve_output(): # Solving Ax=b where A is the identity matrix v.data[:] = 5.0 eqn = Eq(u, v) - petsc = PETScSolve(eqn, target=u) + petsc = petscsolve(eqn, target=u) with switchconfig(language='petsc'): op = Operator(petsc) @@ -568,74 +675,75 @@ def test_solve_output(): assert np.allclose(u.data, v.data) -@skipif('petsc') -def test_essential_bcs(): - """ - Verify that PETScSolve returns the correct output with - essential boundary conditions. - """ - # SubDomains used for essential boundary conditions - # should not overlap. - class SubTop(SubDomain): - name = 'subtop' +class TestEssentialBCs: + @skipif('petsc') + def test_essential_bcs(self): + """ + Verify that `petscsolve` returns the correct output with + essential boundary conditions (`EssentialBC`). + """ + # SubDomains used for essential boundary conditions + # should not overlap. + class SubTop(SubDomain): + name = 'subtop' - def define(self, dimensions): - x, y = dimensions - return {x: x, y: ('right', 1)} - sub1 = SubTop() + def define(self, dimensions): + x, y = dimensions + return {x: x, y: ('right', 1)} + sub1 = SubTop() - class SubBottom(SubDomain): - name = 'subbottom' + class SubBottom(SubDomain): + name = 'subbottom' - def define(self, dimensions): - x, y = dimensions - return {x: x, y: ('left', 1)} - sub2 = SubBottom() + def define(self, dimensions): + x, y = dimensions + return {x: x, y: ('left', 1)} + sub2 = SubBottom() - class SubLeft(SubDomain): - name = 'subleft' + class SubLeft(SubDomain): + name = 'subleft' - def define(self, dimensions): - x, y = dimensions - return {x: ('left', 1), y: ('middle', 1, 1)} - sub3 = SubLeft() + def define(self, dimensions): + x, y = dimensions + return {x: ('left', 1), y: ('middle', 1, 1)} + sub3 = SubLeft() - class SubRight(SubDomain): - name = 'subright' + class SubRight(SubDomain): + name = 'subright' - def define(self, dimensions): - x, y = dimensions - return {x: ('right', 1), y: ('middle', 1, 1)} - sub4 = SubRight() + def define(self, dimensions): + x, y = dimensions + return {x: ('right', 1), y: ('middle', 1, 1)} + sub4 = SubRight() - subdomains = (sub1, sub2, sub3, sub4) - grid = Grid(shape=(11, 11), subdomains=subdomains, dtype=np.float64) + subdomains = (sub1, sub2, sub3, sub4) + grid = Grid(shape=(11, 11), subdomains=subdomains, dtype=np.float64) - u = Function(name='u', grid=grid, space_order=2) - v = Function(name='v', grid=grid, space_order=2) + u = Function(name='u', grid=grid, space_order=2) + v = Function(name='v', grid=grid, space_order=2) - # Solving Ax=b where A is the identity matrix - v.data[:] = 5.0 - eqn = Eq(u, v, subdomain=grid.interior) + # Solving Ax=b where A is the identity matrix + v.data[:] = 5.0 + eqn = Eq(u, v, subdomain=grid.interior) - bcs = [EssentialBC(u, 1., subdomain=sub1)] # top - bcs += [EssentialBC(u, 2., subdomain=sub2)] # bottom - bcs += [EssentialBC(u, 3., subdomain=sub3)] # left - bcs += [EssentialBC(u, 4., subdomain=sub4)] # right + bcs = [EssentialBC(u, 1., subdomain=sub1)] # top + bcs += [EssentialBC(u, 2., subdomain=sub2)] # bottom + bcs += [EssentialBC(u, 3., subdomain=sub3)] # left + bcs += [EssentialBC(u, 4., subdomain=sub4)] # right - petsc = PETScSolve([eqn]+bcs, target=u) + petsc = petscsolve([eqn]+bcs, target=u) - with switchconfig(language='petsc'): - op = Operator(petsc) - op.apply() + with switchconfig(language='petsc'): + op = Operator(petsc) + op.apply() - # Check u is equal to v on the interior - assert np.allclose(u.data[1:-1, 1:-1], v.data[1:-1, 1:-1]) - # Check u satisfies the boundary conditions - assert np.allclose(u.data[1:-1, -1], 1.0) # top - assert np.allclose(u.data[1:-1, 0], 2.0) # bottom - assert np.allclose(u.data[0, 1:-1], 3.0) # left - assert np.allclose(u.data[-1, 1:-1], 4.0) # right + # Check u is equal to v on the interior + assert np.allclose(u.data[1:-1, 1:-1], v.data[1:-1, 1:-1]) + # Check u satisfies the boundary conditions + assert np.allclose(u.data[1:-1, -1], 1.0) # top + assert np.allclose(u.data[1:-1, 0], 2.0) # bottom + assert np.allclose(u.data[0, 1:-1], 3.0) # left + assert np.allclose(u.data[-1, 1:-1], 4.0) # right @skipif('petsc') @@ -668,7 +776,7 @@ def define(self, dimensions): eq1 = Eq(e.laplace + e, f + 2.0) - petsc = PETScSolve([eq1, bc_1, bc_2], target=e) + petsc = petscsolve([eq1, bc_1, bc_2], target=e) jac = petsc.rhs.field_data.jacobian @@ -714,7 +822,7 @@ def define(self, dimensions): eq1 = Eq(e.laplace + e, f + 2.0) - petsc = PETScSolve([eq1, bc_1, bc_2], target=e) + petsc = petscsolve([eq1, bc_1, bc_2], target=e) res = petsc.rhs.field_data.residual @@ -753,7 +861,7 @@ class TestCoupledLinear: def test_coupled_vs_non_coupled(self, eq1, eq2, so): """ Test that solving multiple **uncoupled** equations separately - vs. together with `PETScSolve` yields the same result. + vs. together with `petscsolve` yields the same result. This test is non time-dependent. """ grid = Grid(shape=(11, 11), dtype=np.float64) @@ -769,8 +877,8 @@ def test_coupled_vs_non_coupled(self, eq1, eq2, so): eq2 = eval(eq2) # Non-coupled - petsc1 = PETScSolve(eq1, target=e) - petsc2 = PETScSolve(eq2, target=g) + petsc1 = petscsolve(eq1, target=e) + petsc2 = petscsolve(eq2, target=g) with switchconfig(language='petsc'): op1 = Operator([petsc1, petsc2], opt='noop') @@ -784,7 +892,7 @@ def test_coupled_vs_non_coupled(self, eq1, eq2, so): g.data[:] = 0 # Coupled - petsc3 = PETScSolve({e: [eq1], g: [eq2]}) + petsc3 = petscsolve({e: [eq1], g: [eq2]}) with switchconfig(language='petsc'): op2 = Operator(petsc3, opt='noop') @@ -828,7 +936,7 @@ def test_coupled_structs(self): eq1 = Eq(e + 5, f) eq2 = Eq(g + 10, h) - petsc = PETScSolve({f: [eq1], h: [eq2]}) + petsc = petscsolve({f: [eq1], h: [eq2]}) name = "foo" @@ -873,7 +981,7 @@ def test_coupled_frees(self, n_fields): *solved_funcs, h = functions equations = [Eq(func.laplace, h) for func in solved_funcs] - petsc = PETScSolve({func: [eq] for func, eq in zip(solved_funcs, equations)}) + petsc = petscsolve({func: [eq] for func, eq in zip(solved_funcs, equations)}) with switchconfig(language='petsc'): op = Operator(petsc, opt='noop') @@ -903,9 +1011,9 @@ def test_dmda_dofs(self): eq2 = Eq(f.laplace, h) eq3 = Eq(g.laplace, h) - petsc1 = PETScSolve({e: [eq1]}) - petsc2 = PETScSolve({e: [eq1], f: [eq2]}) - petsc3 = PETScSolve({e: [eq1], f: [eq2], g: [eq3]}) + petsc1 = petscsolve({e: [eq1]}) + petsc2 = petscsolve({e: [eq1], f: [eq2]}) + petsc3 = petscsolve({e: [eq1], f: [eq2], g: [eq3]}) with switchconfig(language='petsc'): op1 = Operator(petsc1, opt='noop') @@ -936,7 +1044,7 @@ def test_mixed_jacobian(self): eq1 = Eq(e.laplace, f) eq2 = Eq(g.laplace, h) - petsc = PETScSolve({e: [eq1], g: [eq2]}) + petsc = petscsolve({e: [eq1], g: [eq2]}) jacobian = petsc.rhs.field_data.jacobian @@ -1048,7 +1156,7 @@ def test_coupling(self, eq1, eq2, j01_matvec, j10_matvec): eq1 = eval(eq1) eq2 = eval(eq2) - petsc = PETScSolve({e: [eq1], g: [eq2]}) + petsc = petscsolve({e: [eq1], g: [eq2]}) jacobian = petsc.rhs.field_data.jacobian @@ -1097,7 +1205,7 @@ def test_jacobian_scaling_1D(self, eq1, eq2, so, scale): eq1 = eval(eq1) eq2 = eval(eq2) - petsc = PETScSolve({e: [eq1], g: [eq2]}) + petsc = petscsolve({e: [eq1], g: [eq2]}) jacobian = petsc.rhs.field_data.jacobian @@ -1147,7 +1255,7 @@ def test_jacobian_scaling_2D(self, eq1, eq2, so, scale): eq1 = eval(eq1) eq2 = eval(eq2) - petsc = PETScSolve({e: [eq1], g: [eq2]}) + petsc = petscsolve({e: [eq1], g: [eq2]}) jacobian = petsc.rhs.field_data.jacobian @@ -1200,7 +1308,7 @@ def test_jacobian_scaling_3D(self, eq1, eq2, so, scale): eq1 = eval(eq1) eq2 = eval(eq2) - petsc = PETScSolve({e: [eq1], g: [eq2]}) + petsc = petscsolve({e: [eq1], g: [eq2]}) jacobian = petsc.rhs.field_data.jacobian @@ -1222,9 +1330,9 @@ def test_residual_bundle(self): eq2 = Eq(f.laplace, h) eq3 = Eq(g.laplace, h) - petsc1 = PETScSolve({e: [eq1]}) - petsc2 = PETScSolve({e: [eq1], f: [eq2]}) - petsc3 = PETScSolve({e: [eq1], f: [eq2], g: [eq3]}) + petsc1 = petscsolve({e: [eq1]}) + petsc2 = petscsolve({e: [eq1], f: [eq2]}) + petsc3 = petscsolve({e: [eq1], f: [eq2], g: [eq3]}) with switchconfig(language='petsc'): op1 = Operator(petsc1, opt='noop', name='op1') @@ -1265,7 +1373,7 @@ def test_residual_callback(self): eq1 = Eq(e.laplace, f) eq2 = Eq(g.laplace, h) - petsc = PETScSolve({e: [eq1], g: [eq2]}) + petsc = petscsolve({e: [eq1], g: [eq2]}) with switchconfig(language='petsc'): op = Operator(petsc) @@ -1312,7 +1420,7 @@ def define(self, dimensions): bc_u = [EssentialBC(u, 0., subdomain=sub1)] bc_v = [EssentialBC(v, 0., subdomain=sub1)] - petsc = PETScSolve({v: [eqn1]+bc_v, u: [eqn2]+bc_u}) + petsc = petscsolve({v: [eqn1]+bc_v, u: [eqn2]+bc_u}) with switchconfig(language='petsc'): op = Operator(petsc) @@ -1390,7 +1498,7 @@ def define(self, dimensions): bcs = [EssentialBC(u, u0, subdomain=sub1)] bcs += [EssentialBC(u, u1, subdomain=sub2)] - petsc = PETScSolve([eqn] + bcs, target=u, solver_parameters={'ksp_rtol': 1e-10}) + petsc = petscsolve([eqn] + bcs, target=u, solver_parameters={'ksp_rtol': 1e-10}) op = Operator(petsc, language='petsc') op.apply() @@ -1414,7 +1522,7 @@ def test_logging(self, log_level): f.data[:] = 5.0 eq = Eq(e.laplace, f) - petsc = PETScSolve(eq, target=e, options_prefix='poisson') + petsc = petscsolve(eq, target=e, options_prefix='poisson') with switchconfig(language='petsc', log_level=log_level): op = Operator(petsc) @@ -1473,8 +1581,8 @@ def test_logging_multiple_solves(self, log_level): eq1 = Eq(g.laplace, e) eq2 = Eq(h, f + 5.0) - solver1 = PETScSolve(eq1, target=g, options_prefix='poisson1') - solver2 = PETScSolve(eq2, target=h, options_prefix='poisson2') + solver1 = petscsolve(eq1, target=g, options_prefix='poisson1') + solver2 = petscsolve(eq2, target=h, options_prefix='poisson2') with switchconfig(language='petsc', log_level=log_level): op = Operator([solver1, solver2]) @@ -1516,8 +1624,8 @@ def test_logging_user_prefixes(self, log_level): pde1 = Eq(e.laplace, f) pde2 = Eq(g.laplace, h) - petsc1 = PETScSolve(pde1, target=e, options_prefix='pde1') - petsc2 = PETScSolve(pde2, target=g, options_prefix='pde2') + petsc1 = petscsolve(pde1, target=e, options_prefix='pde1') + petsc2 = petscsolve(pde2, target=g, options_prefix='pde2') with switchconfig(language='petsc', log_level=log_level): op = Operator([petsc1, petsc2]) @@ -1545,8 +1653,8 @@ def test_logging_default_prefixes(self, log_level): pde1 = Eq(e.laplace, f) pde2 = Eq(g.laplace, h) - petsc1 = PETScSolve(pde1, target=e) - petsc2 = PETScSolve(pde2, target=g) + petsc1 = petscsolve(pde1, target=e) + petsc2 = petscsolve(pde2, target=g) with switchconfig(language='petsc', log_level=log_level): op = Operator([petsc1, petsc2]) @@ -1578,11 +1686,11 @@ def setup_class(self): @skipif('petsc') def test_different_solver_params(self): # Explicitly set the solver parameters - solver1 = PETScSolve( + solver1 = petscsolve( self.eq1, target=self.e, solver_parameters={'ksp_rtol': '1e-10'} ) # Use solver parameter defaults - solver2 = PETScSolve(self.eq2, target=self.g) + solver2 = petscsolve(self.eq2, target=self.g) with switchconfig(language='petsc'): op = Operator([solver1, solver2]) @@ -1598,10 +1706,10 @@ def test_different_solver_params(self): @skipif('petsc') def test_options_prefix(self): - solver1 = PETScSolve(self.eq1, self.e, + solver1 = petscsolve(self.eq1, self.e, solver_parameters={'ksp_rtol': '1e-10'}, options_prefix='poisson1') - solver2 = PETScSolve(self.eq2, self.g, + solver2 = petscsolve(self.eq2, self.g, solver_parameters={'ksp_rtol': '1e-12'}, options_prefix='poisson2') @@ -1625,7 +1733,7 @@ def test_options_no_value(self): Test solver parameters that do not require a value, such as `snes_view` and `ksp_view`. """ - solver = PETScSolve( + solver = petscsolve( self.eq1, target=self.e, solver_parameters={'snes_view': None}, options_prefix='solver1' ) @@ -1645,7 +1753,7 @@ def test_tolerances(self, log_level): 'ksp_divtol': 1e3, 'ksp_max_it': 100 } - solver = PETScSolve( + solver = petscsolve( self.eq1, target=self.e, solver_parameters=params, options_prefix='solver' ) @@ -1668,11 +1776,11 @@ def test_tolerances(self, log_level): @skipif('petsc') def test_clearing_options(self): # Explicitly set the solver parameters - solver1 = PETScSolve( + solver1 = petscsolve( self.eq1, target=self.e, solver_parameters={'ksp_rtol': '1e-10'} ) # Use the solver parameter defaults - solver2 = PETScSolve(self.eq2, target=self.g) + solver2 = petscsolve(self.eq2, target=self.g) with switchconfig(language='petsc'): op = Operator([solver1, solver2]) @@ -1686,11 +1794,11 @@ def test_error_if_same_prefix(self): Test an error is raised if the same options prefix is used for two different solvers within the same Operator. """ - solver1 = PETScSolve( + solver1 = petscsolve( self.eq1, target=self.e, options_prefix='poisson', solver_parameters={'ksp_rtol': '1e-10'} ) - solver2 = PETScSolve( + solver2 = petscsolve( self.eq2, target=self.g, options_prefix='poisson', solver_parameters={'ksp_rtol': '1e-12'} ) @@ -1702,19 +1810,19 @@ def test_error_if_same_prefix(self): @pytest.mark.parametrize('log_level', ['PERF', 'DEBUG']) def test_multiple_operators(self, log_level): """ - Verify that solver parameters are set correctly when multiple Operators - are created with PETScSolve instances sharing the same options_prefix. + Verify that solver parameters are set correctly when multiple `Operator`s + are created with `petscsolve` calls sharing the same `options_prefix`. - Note: Using the same options_prefix within a single Operator is not allowed + Note: Using the same `options_prefix` within a single `Operator` is not allowed (see previous test), but the same prefix can be used across - different Operators (although not advised). + different `Operator`s (although not advised). """ - # Create two PETScSolve instances with the same options_prefix - solver1 = PETScSolve( + # Create two `petscsolve` calls with the same `options_prefix`` + solver1 = petscsolve( self.eq1, target=self.e, options_prefix='poisson', solver_parameters={'ksp_rtol': '1e-10'} ) - solver2 = PETScSolve( + solver2 = petscsolve( self.eq2, target=self.g, options_prefix='poisson', solver_parameters={'ksp_rtol': '1e-12'} ) @@ -1743,7 +1851,7 @@ def test_command_line_priority_tols_1(self, command_line, log_level): prefix = 'd17weqroeg' _, expected = command_line - solver1 = PETScSolve( + solver1 = petscsolve( self.eq1, target=self.e, options_prefix=prefix ) @@ -1762,7 +1870,7 @@ def test_command_line_priority_tols_2(self, command_line, log_level): prefix = 'riabfodkj5' _, expected = command_line - solver1 = PETScSolve( + solver1 = petscsolve( self.eq1, target=self.e, options_prefix=prefix ) @@ -1795,7 +1903,7 @@ def test_command_line_priority_tols3(self, command_line, log_level): 'ksp_max_it': 500 } - solver1 = PETScSolve( + solver1 = petscsolve( self.eq1, target=self.e, solver_parameters=params, options_prefix=prefix @@ -1824,7 +1932,7 @@ def test_command_line_priority_ksp_type(self, command_line, log_level): # see the `command_line` fixture). params = {'ksp_type': 'richardson'} - solver1 = PETScSolve( + solver1 = petscsolve( self.eq1, target=self.e, solver_parameters=params, options_prefix=prefix @@ -1850,7 +1958,7 @@ def test_command_line_priority_ccode(self, command_line): """ prefix = 'qtr2vfvwiu' - solver = PETScSolve( + solver = petscsolve( self.eq1, target=self.e, # Specify a solver parameter that is not set via the # command line (see the `command_line` fixture for this prefix). @@ -1899,21 +2007,21 @@ def test_solveexpr(self): e, f = functions eq = Eq(e.laplace, f) - # Two PETScSolve instances with different options_prefix values + # Two `petscsolve` calls with different `options_prefix` values # should hash differently. - petsc1 = PETScSolve(eq, target=e, options_prefix='poisson1') - petsc2 = PETScSolve(eq, target=e, options_prefix='poisson2') + petsc1 = petscsolve(eq, target=e, options_prefix='poisson1') + petsc2 = petscsolve(eq, target=e, options_prefix='poisson2') assert hash(petsc1.rhs) != hash(petsc2.rhs) assert petsc1.rhs != petsc2.rhs - # Two PETScSolve instances with the same options_prefix but - # different solver parameters should hash differently. - petsc3 = PETScSolve( + # Two `petscsolve` calls with the same `options_prefix` but + # different `solver_parameters` should hash differently. + petsc3 = petscsolve( eq, target=e, solver_parameters={'ksp_type': 'cg'}, options_prefix='poisson3' ) - petsc4 = PETScSolve( + petsc4 = petscsolve( eq, target=e, solver_parameters={'ksp_type': 'richardson'}, options_prefix='poisson3' ) @@ -1922,7 +2030,7 @@ def test_solveexpr(self): class TestGetInfo: """ - Test the `get_info` optional argument to `PETScSolve`. + Test the `get_info` (optional) argument to `petscsolve`. This argument can be used independently of the `log_level` to retrieve specific information about the solve, such as the number of KSP @@ -1945,7 +2053,7 @@ def setup_class(self): @skipif('petsc') def test_get_info(self): get_info = ['kspgetiterationnumber', 'snesgetiterationnumber'] - petsc = PETScSolve( + petsc = petscsolve( self.eq1, target=self.e, options_prefix='pde1', get_info=get_info ) with switchconfig(language='petsc'): @@ -1968,7 +2076,7 @@ def test_get_info_with_logging(self, log_level): Test that `get_info` works correctly when logging is enabled. """ get_info = ['kspgetiterationnumber'] - petsc = PETScSolve( + petsc = petscsolve( self.eq1, target=self.e, options_prefix='pde1', get_info=get_info ) with switchconfig(language='petsc', log_level=log_level): @@ -1991,15 +2099,15 @@ def test_different_solvers(self): Test that `get_info` works correctly when multiple solvers are used within the same Operator. """ - # Create two PETScSolve instances with different get_info arguments + # Create two `petscsolve` calls with different `get_info` arguments get_info_1 = ['kspgetiterationnumber'] get_info_2 = ['snesgetiterationnumber'] - solver1 = PETScSolve( + solver1 = petscsolve( self.eq1, target=self.e, options_prefix='pde1', get_info=get_info_1 ) - solver2 = PETScSolve( + solver2 = petscsolve( self.eq2, target=self.g, options_prefix='pde2', get_info=get_info_2 ) with switchconfig(language='petsc'): @@ -2028,7 +2136,7 @@ def test_case_insensitive(self): """ # Create a list with mixed cases get_info = ['KSPGetIterationNumber', 'snesgetiterationnumber'] - petsc = PETScSolve( + petsc = petscsolve( self.eq1, target=self.e, options_prefix='pde1', get_info=get_info ) with switchconfig(language='petsc'): @@ -2048,10 +2156,10 @@ def test_get_ksp_type(self): a string. """ get_info = ['kspgettype'] - solver1 = PETScSolve( + solver1 = petscsolve( self.eq1, target=self.e, options_prefix='poisson1', get_info=get_info ) - solver2 = PETScSolve( + solver2 = petscsolve( self.eq1, target=self.e, options_prefix='poisson2', solver_parameters={'ksp_type': 'cg'}, get_info=get_info ) @@ -2076,3 +2184,24 @@ def test_get_ksp_type(self): assert entry2.KSPGetType == 'cg' assert entry2['KSPGetType'] == 'cg' assert entry2['kspgettype'] == 'cg' + + +class TestPrinter: + + @skipif('petsc') + def test_petsc_pi(self): + """ + Test that sympy.pi is correctly translated to PETSC_PI in the + generated code. + """ + grid = Grid(shape=(11, 11), dtype=np.float64) + e = Function(name='e', grid=grid) + eq = Eq(e, sp.pi) + + petsc = petscsolve(eq, target=e) + + with switchconfig(language='petsc'): + op = Operator(petsc) + + assert 'PETSC_PI' in str(op.ccode) + assert 'M_PI' not in str(op.ccode) From bf52c5bab33dbc4f22ce4b392178d2d33a0eb5e1 Mon Sep 17 00:00:00 2001 From: ZoeLeibowitz Date: Tue, 23 Sep 2025 13:19:08 +0100 Subject: [PATCH 83/84] misc: Fixes since merge with main --- devito/ir/iet/nodes.py | 16 ++++++++-------- devito/ir/iet/visitors.py | 13 ++++++++++++- tests/test_iet.py | 16 ++++++++-------- 3 files changed, 28 insertions(+), 17 deletions(-) diff --git a/devito/ir/iet/nodes.py b/devito/ir/iet/nodes.py index ed54c8d521..21f88e0350 100644 --- a/devito/ir/iet/nodes.py +++ b/devito/ir/iet/nodes.py @@ -1060,7 +1060,9 @@ class Dereference(ExprStmt, Node): The following cases are supported: * `pointer` is an AbstractFunction, and `pointee` is an Array. - * `pointer` is an AbstractObject, and `pointee` is an Array. + * `pointer` is an AbstractObject, and `pointee` is an Array + - if the `pointer` is a `LocalCompositeObject`, then `pointee` is a + Symbol representing the derefrerenced value. * `pointer` is a Symbol with its _C_ctype deriving from ct._Pointer, and `pointee` is a Symbol representing the dereferenced value. """ @@ -1092,16 +1094,14 @@ def expr_symbols(self): for i in self.pointee.symbolic_shape[1:])) ret.extend(self.pointer.free_symbols) elif self.pointer.is_AbstractObject: - ret.extend([self.pointer, self.pointee.indexed]) + if isinstance(self.pointer, LocalCompositeObject): + ret.extend([self.pointer._C_symbol, self.pointee._C_symbol]) + else: + ret.extend([self.pointer, self.pointee.indexed]) ret.extend(flatten(i.free_symbols for i in self.pointee.symbolic_shape[1:])) else: - # TODO: Might be uneccessary now - if isinstance(self.pointer, LocalCompositeObject) or \ - issubclass(self.pointer._C_ctype, ctypes._Pointer): - ret.extend([self.pointer._C_symbol, self.pointee._C_symbol]) - else: - assert False, f"Unexpected pointer type {type(self.pointer)}" + assert False, f"Unexpected pointer type {type(self.pointer)}" return tuple(filter_ordered(ret)) diff --git a/devito/ir/iet/visitors.py b/devito/ir/iet/visitors.py index 7c350cf51a..5c0bd68ab2 100644 --- a/devito/ir/iet/visitors.py +++ b/devito/ir/iet/visitors.py @@ -515,7 +515,18 @@ def visit_PointerCast(self, o): def visit_Dereference(self, o): a0, a1 = o.functions - if a0.is_AbstractFunction: + # TODO: Temporary fix or fine? — ensures that all objects dereferenced from + # a PETSc struct (e.g., `ctx0`) are handled correctly. + # **Example** + # Need this: struct dataobj *rhs_vec = ctx0->rhs_vec; + # Not this: PetscScalar (* rhs)[rhs_vec->size[1]] = + # (PetscScalar (*)[rhs_vec->size[1]]) ctx0; + # This is the case when a1 is a LocalCompositeObject (i.e a1.is_AbstractObject) + + if a1.is_AbstractObject: + rvalue = f'{a1.name}->{a0._C_name}' + lvalue = self._gen_value(a0, 0) + elif a0.is_AbstractFunction: cstr = self.ccode(a0.indexed._C_typedata) try: diff --git a/tests/test_iet.py b/tests/test_iet.py index 32012ce9c7..6087f624f6 100644 --- a/tests/test_iet.py +++ b/tests/test_iet.py @@ -10,7 +10,7 @@ from devito.ir.iet import ( Call, Callable, Conditional, Definition, DeviceCall, DummyExpr, Iteration, List, KernelLaunch, Lambda, ElementalFunction, CGen, FindSymbols, filter_iterations, - make_efunc, retrieve_iteration_tree, Transformer, Callback, Definition, FindNodes + make_efunc, retrieve_iteration_tree, Transformer, Callback, FindNodes ) from devito.ir import SymbolRegistry from devito.passes.iet.engine import Graph @@ -505,16 +505,16 @@ def test_codegen_quality0(): assert foo1.parameters[0] is a -def test_special_array_definition(): +# def test_special_array_definition(): - class MyArray(Array): - is_extern = True - _data_alignment = False +# class MyArray(Array): +# is_extern = True +# _data_alignment = False - dim = CustomDimension(name='d', symbolic_size=String('')) - a = MyArray(name='a', dimensions=dim, scope='shared', dtype=np.uint8) +# dim = CustomDimension(name='d', symbolic_size=String('')) +# a = MyArray(name='a', dimensions=dim, scope='shared', dtype=np.uint8) - assert str(Definition(a)) == "extern unsigned char a[];" +# assert str(Definition(a)) == "extern unsigned char a[];" def test_list_inline(): From 807396cf8a5cb783da6d989cbc1744f4598f8acb Mon Sep 17 00:00:00 2001 From: ZoeLeibowitz Date: Tue, 23 Sep 2025 18:11:05 +0100 Subject: [PATCH 84/84] tests: Edit test_special_array_definition --- tests/test_iet.py | 18 +++++++++--------- 1 file changed, 9 insertions(+), 9 deletions(-) diff --git a/tests/test_iet.py b/tests/test_iet.py index 6087f624f6..617a79f006 100644 --- a/tests/test_iet.py +++ b/tests/test_iet.py @@ -16,9 +16,9 @@ from devito.passes.iet.engine import Graph from devito.passes.iet.languages.C import CDataManager from devito.symbolics import (Byref, FieldFromComposite, InlineIf, Macro, Class, - String, FLOAT) + FLOAT) from devito.tools import CustomDtype, as_tuple, dtype_to_ctype -from devito.types import CustomDimension, Array, LocalObject, Symbol +from devito.types import CustomDimension, Array, LocalObject, Symbol, Constant @pytest.fixture @@ -505,16 +505,16 @@ def test_codegen_quality0(): assert foo1.parameters[0] is a -# def test_special_array_definition(): +def test_special_array_definition(): -# class MyArray(Array): -# is_extern = True -# _data_alignment = False + class MyArray(Array): + is_extern = True + _data_alignment = False -# dim = CustomDimension(name='d', symbolic_size=String('')) -# a = MyArray(name='a', dimensions=dim, scope='shared', dtype=np.uint8) + dim = CustomDimension(name='d', symbolic_size=Constant(name='size', value=3.0)) + a = MyArray(name='a', dimensions=dim, scope='shared', dtype=np.uint8) -# assert str(Definition(a)) == "extern unsigned char a[];" + assert str(Definition(a)) == "extern unsigned char a[size];" def test_list_inline():