diff --git a/.github/workflows/examples-mpi.yml b/.github/workflows/examples-mpi.yml index a16b2c0196..08c328c61d 100644 --- a/.github/workflows/examples-mpi.yml +++ b/.github/workflows/examples-mpi.yml @@ -17,9 +17,11 @@ on: push: branches: - main + - petsc pull_request: branches: - main + - petsc jobs: build: diff --git a/.github/workflows/examples.yml b/.github/workflows/examples.yml index 6eeb19c921..0265ff0524 100644 --- a/.github/workflows/examples.yml +++ b/.github/workflows/examples.yml @@ -10,9 +10,11 @@ on: push: branches: - main + - petsc pull_request: branches: - main + - petsc jobs: tutorials: diff --git a/.github/workflows/flake8.yml b/.github/workflows/flake8.yml index 8867213b04..be774235b6 100644 --- a/.github/workflows/flake8.yml +++ b/.github/workflows/flake8.yml @@ -10,9 +10,11 @@ on: push: branches: - main + - petsc pull_request: branches: - main + - petsc jobs: flake8: diff --git a/.github/workflows/pytest-core-mpi.yml b/.github/workflows/pytest-core-mpi.yml index 467a146a18..50659fb7ab 100644 --- a/.github/workflows/pytest-core-mpi.yml +++ b/.github/workflows/pytest-core-mpi.yml @@ -10,9 +10,11 @@ on: push: branches: - main + - petsc pull_request: branches: - main + - petsc jobs: test-mpi-basic: diff --git a/.github/workflows/pytest-core-nompi.yml b/.github/workflows/pytest-core-nompi.yml index cf05edb9fd..da0eabd62e 100644 --- a/.github/workflows/pytest-core-nompi.yml +++ b/.github/workflows/pytest-core-nompi.yml @@ -10,9 +10,11 @@ on: push: branches: - main + - petsc pull_request: branches: - main + - petsc jobs: pytest: diff --git a/.github/workflows/pytest-petsc.yml b/.github/workflows/pytest-petsc.yml new file mode 100644 index 0000000000..5851c69594 --- /dev/null +++ b/.github/workflows/pytest-petsc.yml @@ -0,0 +1,97 @@ +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: + - main + - petsc + pull_request: + branches: + - main + - petsc + +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: 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 -f docker/Dockerfile.devito --build-arg base=zoeleibowitz/petsc_image:latest --tag zoeleibowitz/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 + 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 - serial + run: | + ${{ 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 }} mpiexec -n 1 pytest -m parallel --cov --cov-config=.coveragerc --cov-report=xml ${{ env.TESTS }} + + - name: Test examples + run: | + ${{ 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/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/02_biharmonic.py + + - 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/.github/workflows/tutorials.yml b/.github/workflows/tutorials.yml index b77a3c625d..a60aba2e70 100644 --- a/.github/workflows/tutorials.yml +++ b/.github/workflows/tutorials.yml @@ -10,9 +10,11 @@ on: push: branches: - main + - petsc pull_request: branches: - main + - petsc jobs: tutorials: diff --git a/conftest.py b/conftest.py index 666db66eae..8b69d8157d 100644 --- a/conftest.py +++ b/conftest.py @@ -15,6 +15,7 @@ from devito.ir.iet import (FindNodes, FindSymbols, Iteration, ParallelBlock, retrieve_iteration_tree) from devito.tools import as_tuple +from devito.petsc.config import PetscOSError, get_petsc_dir try: from mpi4py import MPI # noqa @@ -34,7 +35,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: @@ -94,6 +95,12 @@ def skipif(items, whole_module=False): if i == 'chkpnt' and Revolver is NoopRevolver: skipit = "pyrevolve not installed" break + if i == 'petsc': + try: + _ = get_petsc_dir() + except PetscOSError: + 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 f226a9ff0d..97301d6839 100644 --- a/devito/core/__init__.py +++ b/devito/core/__init__.py @@ -7,6 +7,7 @@ Cpu64CXXNoopCOperator, Cpu64CXXNoopOmpOperator, Cpu64AdvCXXOperator, Cpu64AdvCXXOmpOperator, Cpu64FsgCXXOperator, Cpu64FsgCXXOmpOperator, + Cpu64NoopPetscOperator, Cpu64AdvPetscOperator, Cpu64CustomCXXOmpOperator, Cpu64CustomCOperator ) from devito.core.intel import ( @@ -45,12 +46,14 @@ operator_registry.add(Cpu64NoopOmpOperator, Cpu64, 'noop', 'Copenmp') 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(Cpu64AdvOmpOperator, Cpu64, 'advanced', 'Copenmp') 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 5471956399..a1c8c23af8 100644 --- a/devito/core/cpu.py +++ b/devito/core/cpu.py @@ -11,13 +11,16 @@ 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 +from devito.petsc.iet.passes import lower_petsc_symbols + __all__ = ['Cpu64NoopCOperator', 'Cpu64NoopOmpOperator', 'Cpu64AdvCOperator', 'Cpu64AdvOmpOperator', 'Cpu64FsgCOperator', 'Cpu64FsgOmpOperator', 'Cpu64CustomOperator', 'Cpu64CustomCXXOperator', 'Cpu64AdvCXXOperator', - 'Cpu64AdvCXXOmpOperator', 'Cpu64FsgCXXOperator', 'Cpu64FsgCXXOmpOperator'] + 'Cpu64AdvCXXOmpOperator', 'Cpu64FsgCXXOperator', 'Cpu64FsgCXXOmpOperator', + 'Cpu64NoopPetscOperator'] class Cpu64OperatorMixin: @@ -142,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 @@ -221,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) @@ -360,6 +369,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 @@ -369,6 +387,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/ir/equations/algorithms.py b/devito/ir/equations/algorithms.py index ce844887aa..49f65a2c5d 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 = {} + # To be updated 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/equations/equation.py b/devito/ir/equations/equation.py index f83dc39c94..f4fe81bcad 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 PetscEq __all__ = ['LoweredEq', 'ClusterizedEq', 'DummyEq', 'OpInc', 'OpMin', 'OpMax', - 'identity_mapper'] + 'identity_mapper', 'OpPetsc'] class IREq(sympy.Eq, Pickable): @@ -102,7 +104,8 @@ def detect(cls, expr): reduction_mapper = { Inc: OpInc, ReduceMax: OpMax, - ReduceMin: OpMin + ReduceMin: OpMin, + PetscEq: OpPetsc } try: return reduction_mapper[type(expr)] @@ -119,6 +122,7 @@ def detect(cls, expr): OpInc = Operation('+') OpMax = Operation('max') OpMin = Operation('min') +OpPetsc = Operation('solve') identity_mapper = { diff --git a/devito/ir/iet/algorithms.py b/devito/ir/iet/algorithms.py index 0b57b876f7..01e3b4976e 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 MetaData +from devito.petsc.iet.nodes 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, MetaData): + 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 5db32afe56..21f88e0350 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,8 @@ 'Increment', 'Return', 'While', 'ListMajor', 'ParallelIteration', 'ParallelBlock', 'Dereference', 'Lambda', 'SyncSpot', 'Pragma', 'DummyExpr', 'BlankLine', 'ParallelTree', 'BusyWait', 'UsingNamespace', - 'Using', 'CallableBody', 'Transfer', 'EmptyList'] + 'Using', 'CallableBody', 'Transfer', 'Callback', 'FixedArgsCallable', + 'EmptyList'] # First-class IET nodes @@ -772,6 +773,15 @@ def defines(self): return self.all_parameters +class FixedArgsCallable(Callable): + + """ + A Callable class that enforces a fixed function signature. + """ + + pass + + class CallableBody(MultiTraversable): """ @@ -1050,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. """ @@ -1082,7 +1094,10 @@ 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: @@ -1151,6 +1166,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 3dcabc5602..5c0bd68ab2 100644 --- a/devito/ir/iet/visitors.py +++ b/devito/ir/iet/visitors.py @@ -26,7 +26,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', 'FindWithin', 'FindSections', @@ -252,7 +252,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: @@ -264,7 +264,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: 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 work + # correctly. + ctype = obj + else: + return None try: return obj._C_typedecl @@ -316,7 +325,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 @@ -506,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: @@ -708,6 +728,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) @@ -790,8 +813,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 @@ -1063,7 +1089,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 @@ -1556,3 +1582,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/logger.py b/devito/logger.py index f3390aef86..58795a64e6 100644 --- a/devito/logger.py +++ b/devito/logger.py @@ -80,9 +80,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: @@ -138,6 +139,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) diff --git a/devito/operator/operator.py b/devito/operator/operator.py index efba54f337..33e3e6a23b 100644 --- a/devito/operator/operator.py +++ b/devito/operator/operator.py @@ -38,7 +38,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'] @@ -199,7 +200,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) @@ -273,6 +274,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 @@ -390,6 +394,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) @@ -487,6 +494,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 @@ -521,7 +531,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'))) @@ -1068,7 +1078,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 f973cfacdd..66c45052d6 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. @@ -209,6 +213,11 @@ def summary(self, args, dtype, 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 @@ -277,7 +286,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 +347,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 if necessary + mapper_func = language_summary_mapper.get(self.language) + if mapper_func: + summary.add_language_summary(self.language, mapper_func(params)) + return summary @@ -366,11 +380,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: @@ -478,6 +492,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 +527,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 +535,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 +557,8 @@ def create_profile(name): 'advisor': AdvisorProfiler } """Profiling levels.""" + + +language_summary_mapper = { + 'petsc': PetscSummary +} 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/definitions.py b/devito/passes/iet/definitions.py index 978c093eed..cb922a3152 100644 --- a/devito/passes/iet/definitions.py +++ b/devito/passes/iet/definitions.py @@ -339,9 +339,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): """ @@ -416,6 +417,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_LocalType), default=float('inf') + )) # maps/unmaps maps = as_list(cbody.maps) + flatten(v.maps) @@ -490,11 +495,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 if isinstance(iet, EntryFunction) and globs: diff --git a/devito/passes/iet/engine.py b/devito/passes/iet/engine.py index 3f002b2c85..26d519e797 100644 --- a/devito/passes/iet/engine.py +++ b/devito/passes/iet/engine.py @@ -4,10 +4,11 @@ import numpy as np from sympy import Mul +from devito.finite_differences.differentiable import Differentiable from devito.ir.iet import ( Call, ExprStmt, Expression, Iteration, SyncSpot, AsyncCallable, FindNodes, - FindSymbols, MapNodes, MetaCall, Transformer, EntryFunction, ThreadCallable, - Uxreplace, derive_parameters + FindSymbols, MapNodes, MetaCall, Transformer, EntryFunction, + FixedArgsCallable, Uxreplace, derive_parameters ) from devito.ir.support import SymbolRegistry from devito.mpi.distributed import MPINeighborhood @@ -24,6 +25,7 @@ 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'] @@ -150,6 +152,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 @@ -518,7 +521,8 @@ def abstract_objects(objects0, sregistry=None): # Precedence rules make it possible to reconstruct objects that depend on # higher priority objects - keys = [Bundle, Array, 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) @@ -553,7 +557,7 @@ def _(i, mapper, sregistry): }) -@abstract_object.register(Array) +@abstract_object.register(ArrayBasic) def _(i, mapper, sregistry): if isinstance(i, Lock): name = sregistry.make_name(prefix='lock') @@ -708,12 +712,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/languages/C.py b/devito/passes/iet/languages/C.py index 7f1c8d8052..6f2850c3a9 100644 --- a/devito/passes/iet/languages/C.py +++ b/devito/passes/iet/languages/C.py @@ -8,6 +8,8 @@ from devito.symbolics import c_complex, c_double_complex from devito.tools import dtype_to_cstr +from devito.petsc.config import petsc_type_mappings + __all__ = ['CBB', 'CDataManager', 'COrchestrator'] @@ -75,3 +77,12 @@ def _print_ComplexPart(self, expr): def _print_Conj(self, expr): # In C, conj is not preceeded by the func_prefix return (f'conj{self.func_literal(expr)}({self._print(expr.args[0])})') + + +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/passes/iet/languages/targets.py b/devito/passes/iet/languages/targets.py index 09fd05b661..95b2417f12 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 CXXDataManager, CXXOrchestrator, CXXPrinter from devito.passes.iet.languages.openmp import (SimdOmpizer, Ompizer, DeviceOmpizer, OmpDataManager, DeviceOmpDataManager, @@ -10,7 +11,7 @@ from devito.passes.iet.instrument import instrument __all__ = ['CTarget', 'OmpTarget', 'COmpTarget', 'DeviceOmpTarget', 'DeviceAccTarget', - 'CXXTarget', 'CXXOmpTarget', 'DeviceCXXOmpTarget'] + 'CXXTarget', 'CXXOmpTarget', 'DeviceCXXOmpTarget', 'PetscTarget'] class Target: @@ -49,6 +50,10 @@ class COmpTarget(Target): Printer = CPrinter +class PetscTarget(CTarget): + Printer = PetscCPrinter + + OmpTarget = COmpTarget diff --git a/devito/petsc/__init__.py b/devito/petsc/__init__.py new file mode 100644 index 0000000000..6aab617e8d --- /dev/null +++ b/devito/petsc/__init__.py @@ -0,0 +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/clusters.py b/devito/petsc/clusters.py new file mode 100644 index 0000000000..00d3e450b7 --- /dev/null +++ b/devito/petsc/clusters.py @@ -0,0 +1,27 @@ +from devito.tools import timed_pass +from devito.petsc.types import SolverMetaData + + +@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, SolverMetaData): + ispace = c.ispace.lift(c.exprs[0].rhs.field_data.space_dimensions) + processed.append(c.rebuild(ispace=ispace)) + else: + processed.append(c) + return processed 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/__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/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 new file mode 100644 index 0000000000..65e2ec2be8 --- /dev/null +++ b/devito/petsc/iet/logging.py @@ -0,0 +1,102 @@ +from functools import cached_property + +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.nodes import petsc_call +from devito.petsc.logging import petsc_return_variable_dict, PetscInfo + + +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) + + if level <= PERF: + funcs = [ + # KSP specific + 'kspgetiterationnumber', + '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: ... + + name = self.sreg.make_name(prefix='petscinfo') + pname = self.sreg.make_name(prefix='petscprofiler') + + self.statstruct = PetscInfo( + name, pname, self.petsc_option_mapper, self.sobjs, + self.section_mapper, self.inject_solve, + self.query_functions + ) + + @cached_property + def petsc_option_mapper(self): + """ + 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): + """ + Generate the PETSc calls that will be injected into the C code to + extract solver statistics. + """ + struct = self.statstruct + calls = [] + for func_name in self.query_functions: + return_variable = petsc_return_variable_dict[func_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(return_variable.name, [input] + by_ref_output) + ) + # TODO: Perform a PetscCIntCast here? + 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/nodes.py b/devito/petsc/iet/nodes.py new file mode 100644 index 0000000000..70508970c3 --- /dev/null +++ b/devito/petsc/iet/nodes.py @@ -0,0 +1,40 @@ +from devito.ir.iet import Expression, Callback, FixedArgsCallable, Call +from devito.ir.equations import OpPetsc + + +class PetscMetaData(Expression): + """ + Base class for general expressions required to run a PETSc solver. + """ + def __init__(self, expr, pragmas=None, operation=OpPetsc): + super().__init__(expr, pragmas=pragmas, operation=operation) + + +class PETScCallable(FixedArgsCallable): + pass + + +class MatShellSetOp(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 + + +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 new file mode 100644 index 0000000000..5154afe43d --- /dev/null +++ b/devito/petsc/iet/passes.py @@ -0,0 +1,302 @@ +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, CallableBody, List, Call, FindNodes, Section, + FindSymbols, DummyExpr, Uxreplace, Dereference +) +from devito.symbolics import Byref, Macro, Null, FieldFromPointer +from devito.types.basic import DataSymbol +import devito.logger + +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, 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 + + +@iet_pass +def lower_petsc(iet, **kwargs): + # Check if `petscsolve` was used + inject_solve_mapper = MapNodes(Iteration, PetscMetaData, + 'groupby').visit(iet) + + if not inject_solve_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']}'" + ) + + data = FindNodes(PetscMetaData).visit(iet) + + if any(filter(lambda i: isinstance(i.expr.rhs, Initialize), data)): + return initialize(iet), core_metadata() + + 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 inject_solve_mapper.values()} + # Assumption is that all solves are on the same `Grid` + if len(unique_grids) > 1: + 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`) + core = make_core_petsc_calls(objs, comm) + + setup = [] + subs = {} + efuncs = {} + + # 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 + 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( + "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, + # executed at the end of the Operator. + clear_options = [] + + for iters, (inject_solve,) in inject_solve_mapper.items(): + + solver = BuildSolver(inject_solve, iters, comm, section_mapper, **kwargs) + + setup.extend(solver.builder.calls) + + # Transform the spatial iteration loop with the calls to execute the solver + subs.update({solver.solve.spatial_body: solver.calls}) + + efuncs.update(solver.callback_builder.efuncs) + + clear_options.extend((petsc_call( + solver.callback_builder._clear_options_efunc.name, [] + ),)) + + populate_matrix_context(efuncs) + iet = Transformer(subs).visit(iet) + 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())} + 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 + argc = DataSymbol(name='argc', dtype=np.int32) + argv = ArgvSymbol(name='argv') + Help = Macro('help') + + 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) + + +class BuildSolver: + """ + 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 `inject_solve`. + """ + 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 = { + 'inject_solve': self.inject_solve, + 'objs': self.objs, + 'iters': self.iters, + 'comm': self.comm, + 'section_mapper': self.section_mapper, + **self.kwargs + } + self.common_kwargs['solver_objs'] = self.type_builder.solver_objs + self.common_kwargs['time_dependence'] = self.time_dependence + self.common_kwargs['callback_builder'] = self.callback_builder + self.common_kwargs['logger'] = self.logger + + @cached_property + def type_builder(self): + return ( + CoupledTypeBuilder(**self.common_kwargs) + if self.coupled else + BaseTypeBuilder(**self.common_kwargs) + ) + + @cached_property + def time_dependence(self): + mapper = self.inject_solve.expr.rhs.time_mapper + time_class = TimeDependent if mapper else TimeIndependent + return time_class(**self.common_kwargs) + + @cached_property + def callback_builder(self): + return CoupledCallbackBuilder(**self.common_kwargs) \ + if self.coupled else BaseCallbackBuilder(**self.common_kwargs) + + @cached_property + def builder(self): + return CoupledBuilder(**self.common_kwargs) \ + if self.coupled else BuilderBase(**self.common_kwargs) + + @cached_property + def solve(self): + return CoupledSolve(**self.common_kwargs) \ + if self.coupled else Solve(**self.common_kwargs) + + @cached_property + def logger(self): + 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) 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/initialize.py b/devito/petsc/initialize.py new file mode 100644 index 0000000000..f225caa4b7 --- /dev/null +++ b/devito/petsc/initialize.py @@ -0,0 +1,60 @@ +import os +import sys +from ctypes import POINTER, cast, c_char +import atexit + +from devito import Operator, switchconfig +from devito.types import Symbol +from devito.types.equation import PetscEq + +from devito.petsc.types import Initialize, Finalize + +global _petsc_initialized +_petsc_initialized = False + + +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 + # from the PETSc module. + 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' + ) + + # 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, 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(clargs), argv=argv_pointer) + atexit.register(op_finalize.apply) + _petsc_initialized = True diff --git a/devito/petsc/logging.py b/devito/petsc/logging.py new file mode 100644 index 0000000000..acbf4cc86f --- /dev/null +++ b/devito/petsc/logging.py @@ -0,0 +1,257 @@ +import os +from collections import namedtuple +from dataclasses import dataclass + +from devito.types import CompositeObject + +from devito.petsc.types import ( + PetscInt, PetscScalar, KSPType, KSPConvergedReason, KSPNormType +) +from devito.petsc.config import petsc_type_to_ctype + + +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 __len__(self): + return len(self._properties) + + def __repr__(self): + return f"PetscEntry({', '.join(f'{k}={v}' for k, v in self.kwargs.items())})" + + +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. + """ + 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.query_functions + )) + self._property_name_map = {} + # Dynamically create a property on this class for each PETSc function + self._add_properties() + + # 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) + + 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. + """ + # 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 + ] + values = [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 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 { + 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', 'petsc_option_mapper', 'sobjs', 'section_mapper', + 'inject_solve', 'query_functions') + + def __init__(self, name, pname, petsc_option_mapper, sobjs, section_mapper, + inject_solve, query_functions): + + self.petsc_option_mapper = petsc_option_mapper + self.sobjs = sobjs + self.section_mapper = section_mapper + self.inject_solve = inject_solve + 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)) + + 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): + section = self.section_mapper.items() + return next((k[0].name for k, v in section if self.inject_solve in v), None) + + @property + def summary_key(self): + return (self.section, self.prefix) + + def __getattr__(self, attr): + 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: tuple + input_params: str + output_param: tuple[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. + +# TODO: To be extended +petsc_return_variable_dict = { + # KSP specific + 'kspgetiterationnumber': PetscReturnVariable( + name='KSPGetIterationNumber', + 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=('snesits',), + ), +} diff --git a/devito/petsc/solve.py b/devito/petsc/solve.py new file mode 100644 index 0000000000..3856392436 --- /dev/null +++ b/devito/petsc/solve.py @@ -0,0 +1,234 @@ +from devito.types.equation import PetscEq +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 +) + + +__all__ = ['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`. + 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_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: + 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. + + 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: + 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_exprs}, + options_prefix, get_info).build_expr() + else: + 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, + get_info=[]): + self.solver_parameters = linear_solver_parameters(solver_parameters) + self.time_mapper = None + self.target_exprs = target_exprs + # 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 = LinearSolverMetaData( + funcs, + solver_parameters=self.solver_parameters, + field_data=field_data, + time_mapper=self.time_mapper, + localinfo=localinfo, + user_prefix=self.user_prefix, + formatted_prefix=self.formatted_prefix, + get_info=self.get_info + ) + return PetscEq(target, linear_solve) + + def linear_solve_args(self): + target, exprs = next(iter(self.target_exprs.items())) + exprs = as_tuple(exprs) + + funcs = get_funcs(exprs) + 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, self.time_mapper) + + field_data = FieldData( + target=target, + jacobian=jacobian, + residual=residual, + initial_guess=initial_guess, + arrays=arrays + ) + + return target, funcs, field_data + + def generate_arrays(self, *targets): + return { + t: { + p: PETScArray( + name=f'{p}_{t.name}', + target=t, + liveness='eager', + localinfo=localinfo + ) + for p in prefixes + } + for t in targets + } + + +class InjectMixedSolve(InjectSolve): + + def linear_solve_args(self): + exprs = [] + for e in self.target_exprs.values(): + exprs.extend(e) + + funcs = get_funcs(exprs) + self.time_mapper = generate_time_mapper(exprs) + + targets = list(self.target_exprs.keys()) + arrays = self.generate_arrays(*targets) + + jacobian = MixedJacobian( + self.target_exprs, arrays, self.time_mapper + ) + + residual = MixedResidual( + self.target_exprs, arrays, self.time_mapper, + jacobian.target_scaler_mapper + ) + + all_data = MultipleFieldData( + targets=targets, + arrays=arrays, + jacobian=jacobian, + residual=residual + ) + + 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') +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..63ea80265b --- /dev/null +++ b/devito/petsc/solver_parameters.py @@ -0,0 +1,42 @@ +import itertools + +from petsctools import flatten_parameters + + +# 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/__init__.py b/devito/petsc/types/__init__.py new file mode 100644 index 0000000000..c40b4acf91 --- /dev/null +++ b/devito/petsc/types/__init__.py @@ -0,0 +1,6 @@ +from .array 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 new file mode 100644 index 0000000000..a56490f3af --- /dev/null +++ b/devito/petsc/types/array.py @@ -0,0 +1,233 @@ +from sympy import Expr + +from functools import cached_property +from ctypes import POINTER, Structure + +from devito.types.utils import DimensionTuple +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 +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 + + __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) + self._coefficients = self._target.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): + target = kwargs['target'] + dimensions = tuple(target.indices[d] for d in 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 grid(self): + return self.target.grid + + @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 staggered(self): + return self.target.staggered + + @property + def is_Staggered(self): + return self.target.staggered is not None + + @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): + 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 = [ + 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): + """ + 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 + + __rkwargs__ = Bundle.__rkwargs__ + ('pname',) + + def __init__(self, *args, pname="Field", **kwargs): + super().__init__(*args, **kwargs) + self._pname = pname + + @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})) + + @property + def symbolic_shape(self): + return self.c0.symbolic_shape + + @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:] + names = tuple(i.target.name for i in self.components) + return PetscComponentAccess( + self.indexed[indices], + component_index, + component_names=names + ) + else: + raise ValueError( + f"Expected {self.ndim} or {self.ndim + 1} indices, " + f"got {len(index)} instead" + ) + + @property + 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(f"Expected Indexed, got `{type(arg)}` instead") + 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): + return self.function._C_ctype diff --git a/devito/petsc/types/equation.py b/devito/petsc/types/equation.py new file mode 100644 index 0000000000..fe9611c1fb --- /dev/null +++ b/devito/petsc/types/equation.py @@ -0,0 +1,45 @@ +from devito.types.equation import Eq + + +__all__ = ['EssentialBC'] + + +class EssentialBC(Eq): + """ + 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 `EssentialBC`s must not overlap. + """ + pass + + +class ZeroRow(EssentialBC): + """ + Equation used to zero all entries, except the diagonal, + of a row in the Jacobian. + + Warnings + -------- + Created and managed directly by Devito, not by users. + """ + pass + + +class ZeroColumn(EssentialBC): + """ + Equation used to zero the column of the Jacobian. + + Warnings + -------- + Created and managed directly by Devito, not by users. + """ + pass diff --git a/devito/petsc/types/macros.py b/devito/petsc/types/macros.py new file mode 100644 index 0000000000..94d9368b5e --- /dev/null +++ b/devito/petsc/types/macros.py @@ -0,0 +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/metadata.py b/devito/petsc/types/metadata.py new file mode 100644 index 0000000000..d36e088a36 --- /dev/null +++ b/devito/petsc/types/metadata.py @@ -0,0 +1,732 @@ +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.symbolics.extraction import separate_eqn, generate_targets, centre_stencil +from devito.types.equation import Eq +from devito.operations.solve import eval_time_derivatives + +from devito.petsc.config 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 + return obj + + +class Initialize(MetaData): + pass + + +class Finalize(MetaData): + pass + + +class GetArgs(MetaData): + pass + + +class SolverMetaData(MetaData): + """ + A symbolic expression passed through the Operator, containing the metadata + needed to execute the PETSc solver. + """ + __rargs__ = ('expr',) + __rkwargs__ = ('solver_parameters', 'field_data', 'time_mapper', + 'localinfo', 'user_prefix', 'formatted_prefix', + 'get_info') + + def __new__(cls, expr, solver_parameters=None, + field_data=None, time_mapper=None, localinfo=None, + user_prefix=None, formatted_prefix=None, + get_info=None, **kwargs): + + with sympy_mutex: + 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 f"{self.__class__.__name__}{self.expr}" + + __str__ = __repr__ + + def _sympystr(self, printer): + return str(self) + + __hash__ = sympy.Basic.__hash__ + + def _hashable_content(self): + return (self.expr, self.formatted_prefix, self.solver_parameters) + + 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): + return self.field_data.grid + + @classmethod + def eval(cls, *args): + return None + + 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 `SolverMetaData`. + Used to interface with PETSc SNES solvers at the IET level. + + 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 residual function F(target) = 0. + 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, + initial_guess=None, arrays=None, **kwargs): + self._target = target + petsc_precision = dtype_mapper[petsc_variables['PETSC_PRECISION']] + if self._target.dtype != petsc_precision: + raise TypeError( + "Your target dtype must match the precision of your " + "PETSc configuration. " + f"Expected {petsc_precision}, but got {self._target.dtype}." + ) + self._jacobian = jacobian + self._residual = residual + self._initial_guess = initial_guess + self._arrays = arrays + + @property + def target(self): + return self._target + + @property + def jacobian(self): + return self._jacobian + + @property + def residual(self): + return self._residual + + @property + def initial_guess(self): + return self._initial_guess + + @property + def arrays(self): + return self._arrays + + @property + def space_dimensions(self): + return self.target.space_dimensions + + @property + def grid(self): + return self.target.grid + + @property + def space_order(self): + return self.target.space_order + + @property + def targets(self): + return as_tuple(self.target) + + +class MultipleFieldData(FieldData): + """ + 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. + + 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 residual function F(targets) = 0. + initial_guess : 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 + self._jacobian = jacobian + self._residual = residual + + @cached_property + def space_dimensions(self): + space_dims = {t.space_dimensions for t in self.targets} + 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` call must have the" + " same space dimensions." + ) + return space_dims.pop() + + @cached_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( + "Multiple `Grid`s detected in `petscsolve`;" + " all targets must share one `Grid`." + ) + return grids.pop() + + @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 + # ... 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` call must have the same space order." + ) + return space_orders.pop() + + @property + 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 + + return [ + m if isinstance(m, EssentialBC) else m._rebuild(rhs=m.rhs * vol) + 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'] + + 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 _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(BaseJacobian): + """ + Represents a Jacobian matrix. + + This Jacobian is defined implicitly via matrix-vector products + 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, exprs, arrays, time_mapper): + super().__init__(arrays=arrays, target=target) + self.exprs = exprs + 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 passed to PETSc via + `MatShellSetOperation(...,MATOP_MULT,(void (*)(void))MatMult)`. + """ + return self._matvecs + + @property + def scdiag(self): + return self._scdiag + + @property + def row_target(self): + return self.target + + @property + def col_target(self): + return self.target + + @property + def zero_memory(self): + return True + + def _build_matvecs(self): + matvecs = [] + for eq in self.exprs: + matvecs.extend( + e for e in self._build_matvec_expr(eq) if e is not None + ) + key = lambda e: not isinstance(e, EssentialBC) + matvecs = tuple(sorted(matvecs, key=key)) + + matvecs = self._scale_non_bcs(matvecs) + scdiag = self._compute_scdiag(matvecs) + matvecs = self._scale_bcs(matvecs, scdiag) + + self._matvecs = matvecs + self._scdiag = scdiag + + +class MixedJacobian(BaseJacobian): + """ + 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_exprs, arrays, time_mapper): + super().__init__(arrays=arrays, target=None) + self.targets = tuple(target_exprs) + self.time_mapper = time_mapper + self._submatrices = [] + self._build_blocks(target_exprs) + + @property + def submatrices(self): + """ + Return a list of all submatrix blocks. + Each block contains metadata about the matrix-vector products. + """ + return self._submatrices + + @property + def n_submatrices(self): + """Return the number of submatrix blocks.""" + return len(self._submatrices) + + @cached_property + def nonzero_submatrices(self): + """Return SubMatrixBlock objects that have non-empty matvecs.""" + return [m for m in self.submatrices if m.matvecs] + + @cached_property + def target_scaler_mapper(self): + """ + Map each row target to the scdiag of its corresponding + diagonal subblock. + """ + mapper = {} + 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_exprs): + """ + Build all SubMatrixBlock objects for the Jacobian. + """ + for i, row_target in enumerate(self.targets): + exprs = target_exprs[row_target] + for j, col_target in enumerate(self.targets): + + 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 + matvecs = tuple( + sorted(matvecs, key=lambda e: not isinstance(e, EssentialBC)) + ) + matvecs = self._scale_non_bcs(matvecs, row_target) + scdiag = self._compute_scdiag(matvecs, col_target) + 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 _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. + """ + for sm in self.submatrices: + if sm.row_idx == row_idx and sm.col_idx == col_idx: + return sm + return None + + def __repr__(self): + summary = ', '.join( + f"{sm.name} (row={sm.row_idx}, col={sm.col_idx})" + for sm in self.submatrices + ) + 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 + + @property + def zero_memory(self): + return False + + def __repr__(self): + return (f"") + + +class Residual: + """ + 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 + equations using Newton-type methods. For linear problems, `SNESKSPONLY` + can be used to perform a single Newton iteration, unifying the + interface for both linear and nonlinear problems. + + 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, exprs, arrays, time_mapper, scdiag): + self.target = target + self.exprs = exprs + self.arrays = arrays + self.time_mapper = time_mapper + self.scdiag = scdiag + self._build_exprs() + + @property + def F_exprs(self): + """ + 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._F_exprs + + @property + def b_exprs(self): + """ + 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._b_exprs + + def _build_exprs(self): + """ + """ + F_exprs = [] + b_exprs = [] + + 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 + b_exprs.extend(self._make_b(e, b)) + + 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] + 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.subs(self.time_mapper), subdomain=eq.subdomain + ) + # Move essential boundary condition to the right-hand side + zero_col = ZeroColumn( + arrays['x'], eq.rhs.subs(self.time_mapper), 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 (Eq(arrays['f'], rhs, subdomain=eq.subdomain),) + + def _make_b(self, expr, b): + b_arr = self.arrays[self.target]['b'] + 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=expr.subdomain),) + + def _scale_bcs(self, eq, scdiag=None): + """ + Scale ZeroRow exprs using scdiag + """ + scdiag = scdiag or self.scdiag + return eq._rebuild(rhs=scdiag * eq.rhs) if isinstance(eq, ZeroRow) else eq + + +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()) + self.arrays = arrays + self.time_mapper = time_mapper + self.scdiag = scdiag + self._build_exprs(target_exprs) + + @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 + + def _build_exprs(self, target_exprs): + residual_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) + ) + ) + self._F_exprs = tuple(sorted( + residual_exprs, key=lambda e: not isinstance(e, EssentialBC) + )) + + 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 + + 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(expr, EssentialBC): + rhs = (self.arrays[target]['x'] - expr.rhs)*self.scdiag[target] + zero_row = ZeroRow( + self.arrays[target]['f'], rhs, subdomain=expr.subdomain + ) + zero_col = ZeroColumn( + self.arrays[target]['x'], expr.rhs, subdomain=expr.subdomain + ) + return (zero_row, zero_col) + + else: + if isinstance(zeroed, (int, float)): + rhs = zeroed * volume + else: + rhs = zeroed.subs(mapper) + rhs = rhs.subs(self.time_mapper)*volume + + return (Eq(self.arrays[target]['f'], rhs, subdomain=expr.subdomain),) + + +class InitialGuess: + """ + Metadata passed to `SolverExpr` to define the initial guess + symbolic expressions, enforcing the initial guess to satisfy essential + boundary conditions. + """ + 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 + def exprs(self): + return self._exprs + + def _build_exprs(self, exprs): + """ + Return a list of initial guess symbolic expressions + that satisfy essential boundary conditions. + """ + self._exprs = tuple([ + eq for eq in + (self._make_initial_guess(e) for e in exprs) + if eq is not None + ]) + + def _make_initial_guess(self, expr): + if isinstance(expr, EssentialBC): + assert expr.lhs == self.target + return Eq( + self.arrays[self.target]['x'], expr.rhs.subs(self.time_mapper), + subdomain=expr.subdomain + ) + else: + return None + + +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 frozendict(zip(targets, array_targets)) 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 new file mode 100644 index 0000000000..8db82be365 --- /dev/null +++ b/devito/petsc/types/object.py @@ -0,0 +1,344 @@ +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, Scalar +) +from devito.symbolics import Byref, cast +from devito.types.basic import DataSymbol, LocalType + +from devito.petsc.iet.nodes import petsc_call + + +class PetscMixin: + @property + def _C_free_priority(self): + if type(self) in FREE_PRIORITY: + return FREE_PRIORITY[type(self)] + else: + return super()._C_free_priority + + +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 + is not destroyed during callback execution. + """ + dtype = CustomDtype('DM') + + +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`. + """ + def __init__(self, *args, dofs=1, **kwargs): + super().__init__(*args, **kwargs) + self._dofs = dofs + + @property + def dofs(self): + return self._dofs + + @property + def _C_free(self): + return petsc_call('DMDestroy', [Byref(self.function)]) + + +DMCast = cast('DM') + + +class CallbackMat(PetscObject): + """ + 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(CallbackMat): + @property + def _C_free(self): + return petsc_call('MatDestroy', [Byref(self.function)]) + + +class CallbackVec(PetscObject): + """ + PETSc vector object (Vec). + """ + dtype = CustomDtype('Vec') + + +class Vec(CallbackVec): + @property + def _C_free(self): + return petsc_call('VecDestroy', [Byref(self.function)]) + + +class PetscMPIInt(PetscObject): + """ + PETSc datatype used to represent `int` parameters + to MPI functions. + """ + dtype = CustomDtype('PetscMPIInt') + + +class PetscInt(PetscObject): + """ + PETSc datatype used to represent `int` parameters + to PETSc functions. + """ + dtype = CustomIntType('PetscInt') + + +class PetscScalar(PetscObject): + dtype = CustomIntType('PetscScalar') + + +class PetscBool(PetscObject): + dtype = CustomDtype('PetscBool') + + +class KSP(PetscObject): + """ + PETSc KSP : Linear Systems Solvers. + Manages Krylov Methods. + """ + dtype = CustomDtype('KSP') + + +class KSPType(PetscObject): + dtype = CustomDtype('KSPType') + + +class KSPNormType(PetscObject): + dtype = CustomDtype('KSPNormType') + + +class CallbackSNES(PetscObject): + """ + PETSc SNES : Non-Linear Systems Solvers. + """ + dtype = CustomDtype('SNES') + + +class SNES(CallbackSNES): + @property + def _C_free(self): + return petsc_call('SNESDestroy', [Byref(self.function)]) + + +class PC(PetscObject): + """ + PETSc object that manages all preconditioners (PC). + """ + dtype = CustomDtype('PC') + + +class KSPConvergedReason(PetscObject): + """ + PETSc object - reason a Krylov method was determined + to have converged or diverged. + """ + dtype = CustomDtype('KSPConvergedReason') + + +class DMDALocalInfo(PetscObject): + """ + PETSc object - C struct containing information + about the local grid. + """ + dtype = CustomDtype('DMDALocalInfo') + + +class PetscErrorCode(PetscObject): + """ + PETSc datatype used to return PETSc error codes. + https://petsc.org/release/manualpages/Sys/PetscErrorCode/ + """ + dtype = CustomDtype('PetscErrorCode') + + +class DummyArg(PetscObject): + """ + A void pointer used to satisfy the function + signature of the `FormFunction` callback. + """ + dtype = CustomDtype('void', modifier='*') + + +class MatReuse(PetscObject): + dtype = CustomDtype('MatReuse') + + +class VecScatter(PetscObject): + dtype = CustomDtype('VecScatter') + + +class StartPtr(PetscObject): + def __init__(self, name, dtype): + super().__init__(name=name) + self.dtype = POINTER(dtype_to_ctype(dtype)) + + +class SingleIS(PetscObject): + dtype = CustomDtype('IS') + + +class PETScStruct(LocalCompositeObject): + + @property + def time_dim_fields(self): + """ + Fields within the struct that are updated during the time loop. + These are not set in the `PopulateUserContext` 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 `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] + + _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'): + 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 PETScArrayObject(PetscMixin, ArrayObject, LocalType): + _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 dim(self): + assert len(self.dimensions) == 1 + return self.dimensions[0] + + @property + def nindices(self): + return self._nindices + + @property + def _C_name(self): + return self.name + + @property + def _mem_stack(self): + return False + + +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): + @property + def _C_ctype(self): + return POINTER(POINTER(c_char)) + + +class NofSubMats(Scalar, LocalType): + pass + + +FREE_PRIORITY = { + PETScArrayObject: 0, + Vec: 1, + Mat: 2, + SNES: 3, + DM: 4, +} 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/devito/symbolics/extended_sympy.py b/devito/symbolics/extended_sympy.py index 4a8d2df206..19a6640f99 100644 --- a/devito/symbolics/extended_sympy.py +++ b/devito/symbolics/extended_sympy.py @@ -12,7 +12,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 @@ -22,7 +22,7 @@ 'MathFunction', 'InlineIf', 'ReservedWord', 'Keyword', 'String', 'Macro', 'Class', 'MacroArgument', 'Deref', 'Namespace', 'Rvalue', 'Null', 'SizeOf', 'rfunc', 'BasicWrapperMixin', 'ValueLimit', - 'VectorAccess'] + 'Mod', 'VectorAccess'] class CondEq(sympy.Eq): @@ -92,9 +92,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) @@ -117,6 +124,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: """ @@ -172,7 +199,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 = [] @@ -257,6 +284,10 @@ def __str__(self): def field(self): return self.call + @property + def dtype(self): + return self.field.dtype + __repr__ = __str__ @@ -550,7 +581,11 @@ class Keyword(ReservedWord): class String(ReservedWord): - pass + + def __str__(self): + return f'"{self.value}"' + + __repr__ = __str__ class Macro(ReservedWord): @@ -881,3 +916,6 @@ def rfunc(func, item, *args): min: Min, max: Max, } + + +Null = Macro('NULL') diff --git a/devito/symbolics/extraction.py b/devito/symbolics/extraction.py new file mode 100644 index 0000000000..45a48c80e8 --- /dev/null +++ b/devito/symbolics/extraction.py @@ -0,0 +1,124 @@ +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. + + Parameters + ---------- + 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 + """ + 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) diff --git a/devito/symbolics/inspection.py b/devito/symbolics/inspection.py index 0ff7fcf6ba..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 + 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.free_symbols: - try: - dtypes.add(i.dtype) - except AttributeError: - pass + inspect_args(expr, dtypes) dtype = infer_dtype(dtypes) diff --git a/devito/tools/dtypes_lowering.py b/devito/tools/dtypes_lowering.py index 7f0ba56911..c92af7a36c 100644 --- a/devito/tools/dtypes_lowering.py +++ b/devito/tools/dtypes_lowering.py @@ -16,7 +16,7 @@ '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', 'extract_dtype', 'CustomDtype', - 'mpi4py_mapper'] + 'mpi4py_mapper', 'CustomIntType'] # *** Custom np.dtypes @@ -127,6 +127,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 @@ -319,6 +324,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 c2755a4c42..1d6a40921e 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 + @property def c0(self): # ArrayBasic can be used as a base class for tensorial objects (that is, @@ -526,8 +541,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): @@ -582,19 +599,22 @@ def component_indices(self): class ComponentAccess(Expr, Pickable): - _component_names = ('x', 'y', 'z', 'w') + _default_component_names = ('x', 'y', 'z', 'w') __rargs__ = ('arg',) - __rkwargs__ = ('index',) + __rkwargs__ = ('index', 'component_names') - 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 @@ -623,6 +643,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/devito/types/equation.py b/devito/types/equation.py index b3353e2cb4..5f027fef62 100644 --- a/devito/types/equation.py +++ b/devito/types/equation.py @@ -237,3 +237,7 @@ class ReduceMax(Reduction): class ReduceMin(Reduction): pass + + +class PetscEq(Eq): + pass diff --git a/devito/types/grid.py b/devito/types/grid.py index 6107799ce0..b99fd01c4f 100644 --- a/devito/types/grid.py +++ b/devito/types/grid.py @@ -281,10 +281,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): diff --git a/devito/types/object.py b/devito/types/object.py index 637e19dea0..3a5dead78a 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'] + +__all__ = ['Object', 'LocalObject', 'CompositeObject', 'LocalCompositeObject'] 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): @@ -236,6 +237,41 @@ 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 LocalCompositeObject(CompositeObject, LocalType): + + """ + Object with composite type (e.g., a C struct) defined in C. + """ + + __rargs__ = ('name', 'pname', 'fields') + __rkwargs__ = ('modifier', 'liveness') + + 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 + self.modifier = modifier + assert liveness in ['eager', 'lazy'] + self._liveness = liveness + self._fields = fields + + @property + 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/docker/Dockerfile.devito b/docker/Dockerfile.devito index 51ef3140f0..d33810058c 100644 --- a/docker/Dockerfile.devito +++ b/docker/Dockerfile.devito @@ -102,4 +102,3 @@ USER app EXPOSE 8888 ENTRYPOINT ["/docker-entrypoint.sh"] CMD ["/jupyter"] - diff --git a/docker/Dockerfile.petsc b/docker/Dockerfile.petsc new file mode 100644 index 0000000000..056c860688 --- /dev/null +++ b/docker/Dockerfile.petsc @@ -0,0 +1,28 @@ +############################################################## +# 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 + +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 && \ + mkdir -p /opt/petsc && \ + cd /opt/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) \ + --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" diff --git a/examples/petsc/Poisson/01_poisson.py b/examples/petsc/Poisson/01_poisson.py new file mode 100644 index 0000000000..318e77f000 --- /dev/null +++ b/examples/petsc/Poisson/01_poisson.py @@ -0,0 +1,115 @@ +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: ('middle', 1, 1)} + + +class SubRight(SubDomain): + name = 'subright' + + def define(self, dimensions): + x, y = dimensions + return {x: ('right', 1), y: ('middle', 1, 1)} + + +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) + bc = Function(name='bc', 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)) + + bc.data[:] = 0.0 + + # # Create boundary condition expressions using subdomains + 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}) + + 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..7ac621d1db --- /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: ('middle', 1, 1), y: ('right', 1)} + + +class SubBottom(SubDomain): + name = 'subbottom' + + def define(self, dimensions): + x, y = dimensions + return {x: ('middle', 1, 1), 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/Poisson/03_poisson.py b/examples/petsc/Poisson/03_poisson.py new file mode 100644 index 0000000000..dd72f265ff --- /dev/null +++ b/examples/petsc/Poisson/03_poisson.py @@ -0,0 +1,99 @@ +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 + +# 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..f5f618e5b2 --- /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: ('middle', 1, 1), y: ('right', 1)} + + +class SubBottom(SubDomain): + name = 'subbottom' + + def define(self, dimensions): + x, y = dimensions + return {x: ('middle', 1, 1), 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 diff --git a/examples/petsc/cfd/01_navierstokes.py b/examples/petsc/cfd/01_navierstokes.py new file mode 100644 index 0000000000..13b9e5d450 --- /dev/null +++ b/examples/petsc/cfd/01_navierstokes.py @@ -0,0 +1,299 @@ +import os +import numpy as np + +from devito import (Grid, TimeFunction, Constant, Eq, + Operator, 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 = TimeFunction(name='pn1', grid=grid, space_order=2, dtype=np.float64) + +eq_pn1 = Eq(pn1.forward.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.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.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) + +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.forward.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) + +# 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/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..76b0aac957 --- /dev/null +++ b/examples/petsc/petsc_test.py @@ -0,0 +1,31 @@ +import os +import numpy as np + +from devito import (Grid, Function, Eq, Operator, configuration, + switchconfig) +from devito.petsc import petscsolve +from devito.petsc.initialize import PetscInitialize +configuration['compiler'] = 'custom' +os.environ['CC'] = 'mpicc' + +PetscInitialize() + +nx = 81 +ny = 81 + +grid = Grid(shape=(nx, ny), extent=(2., 2.), dtype=np.float64) + +u = Function(name='u', grid=grid, dtype=np.float64, space_order=2) +v = Function(name='v', grid=grid, dtype=np.float64, space_order=2) + +v.data[:] = 5.0 + +eq = Eq(v, u.laplace, subdomain=grid.interior) + +petsc = petscsolve([eq], u) + +with switchconfig(language='petsc'): + op = Operator(petsc) + op.apply() + +print(op.ccode) diff --git a/examples/petsc/random/01_helmholtz.py b/examples/petsc/random/01_helmholtz.py new file mode 100644 index 0000000000..dc498a0dae --- /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/random/02_biharmonic.py b/examples/petsc/random/02_biharmonic.py new file mode 100644 index 0000000000..635c4ff42b --- /dev/null +++ b/examples/petsc/random/02_biharmonic.py @@ -0,0 +1,155 @@ +# 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: ('middle', 1, 1), y: ('right', 1)} + + +class SubBottom(SubDomain): + name = 'subbottom' + + def define(self, dimensions): + x, y = dimensions + return {x: ('middle', 1, 1), 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] +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 + 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) + 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_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_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) +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 diff --git a/examples/petsc/seismic/01_staggered_acoustic.py b/examples/petsc/seismic/01_staggered_acoustic.py new file mode 100644 index 0000000000..2352083236 --- /dev/null +++ b/examples/petsc/seismic/01_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) 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/requirements.txt b/requirements.txt index fe9e5f1a29..6075a1feba 100644 --- a/requirements.txt +++ b/requirements.txt @@ -7,4 +7,5 @@ cgen>=2020.1,<2026 codepy>=2019.1,<2025 multidict<6.3 anytree>=2.4.3,<=2.13.0 -packaging<25.1 \ No newline at end of file +packaging<25.1 +petsctools<=2025.1 diff --git a/tests/test_iet.py b/tests/test_iet.py index 4b3961d9b4..617a79f006 100644 --- a/tests/test_iet.py +++ b/tests/test_iet.py @@ -10,15 +10,15 @@ 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 + make_efunc, retrieve_iteration_tree, Transformer, Callback, FindNodes ) from devito.ir import SymbolRegistry 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 @@ -129,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'))) @@ -484,10 +511,10 @@ class MyArray(Array): is_extern = True _data_alignment = False - dim = CustomDimension(name='d', symbolic_size=String('')) + 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(): diff --git a/tests/test_petsc.py b/tests/test_petsc.py new file mode 100644 index 0000000000..6926ae76cd --- /dev/null +++ b/tests/test_petsc.py @@ -0,0 +1,2207 @@ +import pytest + +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, sin) +from devito.operator.profiling import PerformanceSummary +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 +from devito.petsc.types import (DM, Mat, Vec, PetscMPIInt, KSP, + PC, KSPConvergedReason, PETScArray, + 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(command_line): + argv, _ = command_line + # TODO: Temporary workaround until PETSc is automatically + # initialized + configuration['compiler'] = 'custom' + os.environ['CC'] = 'mpicc' + PetscInitialize(argv) + + +@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(): + """ + Test C++ support for PETSc LocalObjects. + """ + lo0 = DM('da', stencil_width=1) + lo1 = Mat('A') + lo2 = Vec('x') + lo3 = PetscMPIInt('size') + lo4 = KSP('ksp') + lo5 = PC('pc') + lo6 = KSPConvergedReason('reason') + + iet = Call('foo', [lo0, lo1, lo2, lo3, lo4, lo5, lo6]) + iet = ElementalFunction('foo', iet, parameters=()) + + dm = CDataManager(sregistry=None) + iet = CDataManager.place_definitions.__wrapped__(dm, iet)[0] + + assert 'DM da;' in str(iet) + assert 'Mat A;' in str(iet) + assert 'Vec x;' in str(iet) + assert 'PetscMPIInt size;' in str(iet) + assert 'KSP ksp;' in str(iet) + assert 'PC pc;' in str(iet) + assert 'KSPConvergedReason reason;' in str(iet) + + +@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), dtype=np.float64) + + 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(language='petsc'): + op = Operator(petsc, opt='noop') + + callable_roots = [meta_call.root for meta_call in op._func_table.values()] + + matvec_efunc = [root for root in callable_roots if root.name == 'MatMult0'] + + b_efunc = [root for root in callable_roots if root.name == 'FormRHS0'] + + 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' + ' - 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)*ctx0->h_x*ctx0->h_y' + ) + + 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 + 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_efunc[0].parameters) == 3 + + +@skipif('petsc') +def test_multiple_petsc_solves(): + """ + Test multiple `petscsolve` calls, passed to a single `Operator`. + """ + 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) + + 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, 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, SetPetscOptions + # and ClearPetscOptions per solve. + # TODO: Some efuncs are not reused where reuse is possible — investigate. + assert len(callable_roots) == 12 + + +@skipif('petsc') +def test_petsc_cast(): + """ + Test casting of PETScArray. + """ + 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) + 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(language='petsc'): + op1 = Operator(petsc1) + op2 = Operator(petsc2) + op3 = Operator(petsc3) + + 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_dmda_create(): + + 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) + 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(language='petsc'): + op1 = Operator(petsc1, opt='noop') + op2 = Operator(petsc2, opt='noop') + op3 = Operator(petsc3, opt='noop') + + assert 'PetscCall(DMDACreate1d(PETSC_COMM_WORLD,DM_BOUNDARY_GHOSTED,' + \ + '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));' \ + 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) + + +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) + + name = "foo" + + 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) + + 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_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) + + with switchconfig(language='petsc'): + op1 = Operator(petsc, name="foo1") + op2 = Operator(petsc, name="foo2") + + 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') +def test_callback_arguments(): + """ + Test the arguments of each callback function. + """ + 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) + + eqn1 = Eq(f1.laplace, g1) + + petsc1 = petscsolve(eqn1, f1) + + with switchconfig(language='petsc'): + op = Operator(petsc1) + + 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, X, Y)' + assert str(ff.parameters) == '(snes, X, F, dummy)' + + +@skipif('petsc') +def test_apply(): + + grid = Grid(shape=(13, 13), 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, dtype=np.float64) + + eqn = Eq(pn.laplace*mu, rhs, subdomain=grid.interior) + + petsc = petscsolve(eqn, pn) + + with switchconfig(language='petsc'): + # Build the op + op = Operator(petsc) + + # Check the Operator runs without errors + op.apply() + + # Verify that users can override `mu` + mu_new = Constant(name='mu_new', value=4.0, dtype=np.float64) + op.apply(mu=mu_new) + + +@skipif('petsc') +def test_petsc_frees(): + + 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.laplace, g) + petsc = petscsolve(eqn, f) + + with switchconfig(language='petsc'): + op = Operator(petsc) + + 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));' + + +@skipif('petsc') +def test_calls_to_callbacks(): + + 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.laplace, g) + petsc = petscsolve(eqn, f) + + with switchconfig(language='petsc'): + op = Operator(petsc) + + ccode = str(op.ccode) + + assert '(void (*)(void))MatMult0' in ccode + assert 'PetscCall(SNESSetFunction(snes0,NULL,FormFunction0,(void*)(da0)));' 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), 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(language='petsc'): + op1 = Operator(petsc1) + + # Verify the case with modulo time stepping + 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, save=5) + eq2 = Eq(u2.dt, u2.laplace, subdomain=grid.interior) + petsc2 = petscsolve(eq2, u2.forward) + + with switchconfig(language='petsc'): + op2 = Operator(petsc2) + + assert ('PetscScalar * u2_ptr0 = (time + 1)*localsize0 + ' + '(PetscScalar*)(u2_vec->data);') in str(op2) + + +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) + + 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) + + with switchconfig(language='petsc'): + op2 = Operator(petsc2) + op2.apply(time_M=3) + body2 = str(op2.body) + rhs2 = str(op2._func_table['FormRHS0'].root.ccode) + + assert 'ctx0.time = time' in body2 + assert 'ctx0->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(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) + + 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') + @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) + + 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. forming 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 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() + + 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: ('middle', 1, 1)} + sub3 = SubLeft() + + class SubRight(SubDomain): + name = '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) + + 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) + + 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 + + +@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.field_data.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) + + +@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.field_data.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. + # 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. + + @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, 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=eval(so)) + for n in ['e', 'f', 'g', 'h']] + e, f, g, h = functions + + f.data[:] = 5. + h.data[:] = 5. + + eq1 = eval(eq1) + eq2 = eval(eq2) + + # Non-coupled + petsc1 = petscsolve(eq1, target=e) + petsc2 = petscsolve(eq2, target=g) + + with switchconfig(language='petsc'): + op1 = Operator([petsc1, petsc2], opt='noop') + op1.apply() + + enorm1 = norm(e) + gnorm1 = norm(g) + + # Reset + e.data[:] = 0 + g.data[:] = 0 + + # Coupled + petsc3 = petscsolve({e: [eq1], g: [eq2]}) + + with switchconfig(language='petsc'): + op2 = Operator(petsc3, opt='noop') + op2.apply() + + enorm2 = norm(e) + gnorm2 = norm(g) + + print('enorm1:', enorm1) + print('enorm2:', enorm2) + 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()] + + # Solving for multiple fields within the same matrix system requires + # less callback functions than solving them separately. + # 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 + field1 = petsc2.rhs.field_data + field2 = petsc3.rhs.field_data + + assert isinstance(field0, FieldData) + assert isinstance(field1, FieldData) + assert isinstance(field2, MultipleFieldData) + + @skipif('petsc') + def test_coupled_structs(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 + + eq1 = Eq(e + 5, f) + eq2 = Eq(g + 10, h) + + petsc = petscsolve({f: [eq1], h: [eq2]}) + + name = "foo" + + 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, 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 + + # The public struct Field0 only appears in the header file + 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, n_fields): + grid = Grid(shape=(11, 11), dtype=np.float64) + + functions = [Function(name=f'u{i}', grid=grid, space_order=2) + for i in range(n_fields + 1)] + *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)}) + + with switchconfig(language='petsc'): + op = Operator(petsc, opt='noop') + + frees = op.body.frees + + # 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): + 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) + + 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') + 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_mixed_jacobian(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 + + eq1 = Eq(e.laplace, f) + eq2 = Eq(g.laplace, h) + + petsc = petscsolve({e: [eq1], g: [eq2]}) + + jacobian = petsc.rhs.field_data.jacobian + + 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 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 + + # Technically a non-coupled problem, so the only non-zero submatrices + # should be the diagonal ones i.e J00 and 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),' \ + + ' 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),' \ + + ' h_x*h_y*(Derivative(x_g(x, y), (x, 2)) + Derivative(x_g(x, y), (y, 2))))' + + # Check the col_targets + assert j00.col_target == e + assert j01.col_target == g + 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) + + e = Function(name='e', grid=grid, space_order=2) + g = Function(name='g', grid=grid, space_order=2) + + eq1 = eval(eq1) + eq2 = eval(eq2) + + petsc = petscsolve({e: [eq1], g: [eq2]}) + + jacobian = petsc.rhs.field_data.jacobian + + j01 = jacobian.get_submatrix(0, 1) + j10 = jacobian.get_submatrix(1, 0) + + 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.rhs.field_data.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.rhs.field_data.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.rhs.field_data.jacobian + + j00 = jacobian.get_submatrix(0, 0) + j11 = jacobian.get_submatrix(1, 1) + + assert str(j00.scdiag) == scale + assert str(j11.scdiag) == scale + + @skipif('petsc') + def test_residual_bundle(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 + + 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(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_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): + """ + 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()) == 9 + # 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 + + +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=[2, 4, 8]) + def test_laplacian_1d(self, nx, unorm, mode): + """ + """ + configuration['compiler'] = 'custom' + os.environ['CC'] = 'mpicc' + PetscInitialize() + + 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: (self.side, 1)} + + 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) + + 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)) + + # Create boundary condition expressions using subdomains + bcs = [EssentialBC(u, u0, subdomain=sub1)] + bcs += [EssentialBC(u, u1, subdomain=sub2)] + + petsc = petscsolve([eqn] + bcs, target=u, solver_parameters={'ksp_rtol': 1e-10}) + + op = Operator(petsc, language='petsc') + 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) + + +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 entry0.SNESGetIterationNumber == 1 + + snesits0 = petsc_summary.SNESGetIterationNumber + snesits1 = petsc_summary['SNESGetIterationNumber'] + # Check case insensitive key access + snesits2 = petsc_summary['snesgetiterationnumber'] + snesits3 = petsc_summary['SNESgetiterationNumber'] + + assert snesits0 == snesits1 == snesits2 == snesits3 + + 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') + @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) + 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=log_level): + 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 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 + + @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 `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 + (see previous test), but the same prefix can be used across + different `Operator`s (although not advised). + """ + # 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( + 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` 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') + + assert hash(petsc1.rhs) != hash(petsc2.rhs) + assert petsc1.rhs != petsc2.rhs + + # 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( + 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` calls 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' + + +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) diff --git a/tests/test_symbolics.py b/tests/test_symbolics.py index d227853b4d..770f9a2fe8 100644 --- a/tests/test_symbolics.py +++ b/tests/test_symbolics.py @@ -15,11 +15,13 @@ retrieve_functions, retrieve_indexed, evalrel, CallFromPointer, Cast, # noqa DefFunction, FieldFromPointer, INT, FieldFromComposite, IntDiv, Namespace, Rvalue, ReservedWord, ListInitializer, uxreplace, pow_to_mul, - retrieve_derivatives, BaseCast, SizeOf, VectorAccess + retrieve_derivatives, BaseCast, SizeOf, VectorAccess, separate_eqn, + centre_stencil, sympy_dtype ) from devito.tools import as_tuple, CustomDtype 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 @@ -264,6 +266,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') @@ -308,7 +321,10 @@ 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' + # 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' def test_integer_abs(): @@ -976,6 +992,244 @@ def test_print_div(): 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 + + def test_customdtype_complex(): """ Test that `CustomDtype` doesn't brak is_imag