Coupled (#41)

* compiler/dsl: Add machinery to support coupled solvers

Author: ZoeLeibowitz

.github/workflows/pytest-petsc.yml

@@ -64,7 +64,7 @@ jobs:
 
     - name: Test with pytest
       run: |
-        ${{ env.RUN_CMD }} pytest --cov --cov-config=.coveragerc --cov-report=xml ${{ env.TESTS }}
+        ${{ env.RUN_CMD }} mpiexec -n 1 pytest --cov --cov-config=.coveragerc --cov-report=xml ${{ env.TESTS }}
 
     - name: Upload coverage to Codecov
       if: "!contains(matrix.name, 'docker')"

devito/ir/equations/algorithms.py

@@ -191,7 +191,8 @@ def concretize_subdims(exprs, **kwargs):
     """
     sregistry = kwargs.get('sregistry')
 
-    mapper = {}
+    # Update based on changes in #2509
+    mapper = kwargs.get('concretize_mapper', {})
     rebuilt = {}  # Rebuilt implicit dims etc which are shared between dimensions
 
     _concretize_subdims(exprs, mapper, rebuilt, sregistry)

devito/ir/iet/nodes.py

@@ -1046,7 +1046,7 @@ class Dereference(ExprStmt, Node):
     The following cases are supported:
 
         * `pointer` is a PointerArray or TempFunction, and `pointee` is an Array.
-        * `pointer` is an ArrayObject or CCompositeObject representing a pointer
+        * `pointer` is an ArrayObject or LocalCompositeObject representing a pointer
            to a C struct, and `pointee` is a field in `pointer`.
         * `pointer` is a Symbol with its _C_ctype deriving from ct._Pointer, and
           `pointee` is a Symbol representing the dereferenced value.
@@ -1070,8 +1070,7 @@ def functions(self):
     def expr_symbols(self):
         ret = []
         if self.pointer.is_Symbol:
-            assert (isinstance(self.pointer, LocalCompositeObject) or
-                    issubclass(self.pointer._C_ctype, ctypes._Pointer)), \
+            assert issubclass(self.pointer._C_ctype, ctypes._Pointer), \
                    "Scalar dereference must have a pointer ctype"
             ret.extend([self.pointer._C_symbol, self.pointee._C_symbol])
         elif self.pointer.is_PointerArray or self.pointer.is_TempFunction:
@@ -1080,7 +1079,9 @@ def expr_symbols(self):
                                for i in self.pointee.symbolic_shape[1:]))
             ret.extend(self.pointer.free_symbols)
         else:
-            ret.extend([self.pointer.indexed, self.pointee._C_symbol])
+            assert (isinstance(self.pointer, LocalCompositeObject) or
+                    issubclass(self.pointer._C_ctype, ctypes._Pointer))
+            ret.extend([self.pointer._C_symbol, self.pointee._C_symbol])
         return tuple(filter_ordered(ret))
 
     @property

devito/ir/iet/visitors.py

@@ -25,7 +25,7 @@
                           c_restrict_void_p, sorted_priority)
 from devito.types.basic import AbstractFunction, AbstractSymbol, Basic
 from devito.types import (ArrayObject, CompositeObject, Dimension, Pointer,
-                          IndexedData, DeviceMap)
+                          IndexedData, DeviceMap, LocalCompositeObject)
 
 
 __all__ = ['FindApplications', 'FindNodes', 'FindSections', 'FindSymbols',
@@ -201,7 +201,7 @@ def _restrict_keyword(self):
 
     def _gen_struct_decl(self, obj, masked=()):
         """
-        Convert ctypes.Struct -> cgen.Structure.
+        Convert ctypes.Struct and LocalCompositeObject -> cgen.Structure.
         """
         ctype = obj._C_ctype
         try:
@@ -213,7 +213,16 @@ def _gen_struct_decl(self, obj, masked=()):
                 return None
         except TypeError:
             # E.g., `ctype` is of type `dtypes_lowering.CustomDtype`
-            return None
+            if isinstance(obj, LocalCompositeObject):
+                # TODO: Potentially re-evaluate: Setting ctype to obj allows
+                # _gen_struct_decl to generate a cgen.Structure from a
+                # LocalCompositeObject, where obj._C_ctype is a CustomDtype.
+                # LocalCompositeObject has a __fields__ property,
+                # which allows the subsequent code in this function to function
+                # correctly.
+                ctype = obj
+            else:
+                return None
 
         try:
             return obj._C_typedecl
@@ -718,8 +727,11 @@ def _operator_typedecls(self, o, mode='all'):
         for i in o._func_table.values():
             if not i.local:
                 continue
-            typedecls.extend([self._gen_struct_decl(j) for j in i.root.parameters
-                              if xfilter(j)])
+            typedecls.extend([
+                self._gen_struct_decl(j)
+                for j in FindSymbols().visit(i.root)
+                if xfilter(j)
+            ])
         typedecls = filter_sorted(typedecls, key=lambda i: i.tpname)
 
         return typedecls

devito/passes/iet/definitions.py

@@ -256,9 +256,10 @@ def _alloc_object_array_on_low_lat_mem(self, site, obj, storage):
         """
         Allocate an Array of Objects in the low latency memory.
         """
+        frees = getattr(obj, '_C_free', None)
         decl = Definition(obj)
 
-        storage.update(obj, site, allocs=decl)
+        storage.update(obj, site, allocs=decl, frees=frees)
 
     def _alloc_pointed_array_on_high_bw_mem(self, site, obj, storage):
         """
@@ -335,7 +336,7 @@ def _inject_definitions(self, iet, storage):
             frees.extend(as_list(cbody.frees) + flatten(v.frees))
             frees = sorted(frees, key=lambda x: min(
                 (obj._C_free_priority for obj in FindSymbols().visit(x)
-                 if obj.is_LocalObject), default=float('inf')
+                 if obj.is_LocalType), default=float('inf')
             ))
 
             # maps/unmaps

devito/petsc/clusters.py

@@ -20,7 +20,7 @@ def petsc_lift(clusters):
     processed = []
     for c in clusters:
         if isinstance(c.exprs[0].rhs, LinearSolveExpr):
-            ispace = c.ispace.lift(c.exprs[0].rhs.target.space_dimensions)
+            ispace = c.ispace.lift(c.exprs[0].rhs.fielddata.space_dimensions)
             processed.append(c.rebuild(ispace=ispace))
         else:
             processed.append(c)

devito/petsc/iet/nodes.py

@@ -14,7 +14,7 @@ class PETScCallable(FixedArgsCallable):
     pass
 
 
-class MatVecCallback(Callback):
+class MatShellSetOp(Callback):
     @property
     def callback_form(self):
         param_types_str = ', '.join([str(t) for t in self.param_types])

devito/petsc/iet/passes.py

@@ -1,18 +1,27 @@
 import cgen as c
 import numpy as np
+from functools import cached_property
 
 from devito.passes.iet.engine import iet_pass
 from devito.ir.iet import (Transformer, MapNodes, Iteration, BlankLine,
-                           FindNodes, Call, CallableBody)
-from devito.symbolics import Byref, Macro
+                           DummyExpr, CallableBody, List, Call, Callable,
+                           FindNodes)
+from devito.symbolics import Byref, Macro, FieldFromPointer
+from devito.types import Symbol, Scalar
 from devito.types.basic import DataSymbol
-from devito.petsc.types import (PetscMPIInt, PetscErrorCode, Initialize,
-                                Finalize, ArgvSymbol)
+from devito.tools import frozendict
+from devito.petsc.types import (PetscMPIInt, PetscErrorCode, MultipleFieldData,
+                                PointerIS, Mat, LocalVec, GlobalVec, CallbackMat, SNES,
+                                DummyArg, PetscInt, PointerDM, PointerMat, MatReuse,
+                                CallbackPointerIS, CallbackPointerDM, JacobianStruct,
+                                SubMatrixStruct, Initialize, Finalize, ArgvSymbol)
 from devito.petsc.types.macros import petsc_func_begin_user
 from devito.petsc.iet.nodes import PetscMetaData
 from devito.petsc.utils import core_metadata
-from devito.petsc.iet.routines import (CallbackBuilder, BaseObjectBuilder, BaseSetup,
-                                       Solver, TimeDependent, NonTimeDependent)
+from devito.petsc.iet.routines import (CBBuilder, CCBBuilder, BaseObjectBuilder,
+                                       CoupledObjectBuilder, BaseSetup, CoupledSetup,
+                                       Solver, CoupledSolver, TimeDependent,
+                                       NonTimeDependent)
 from devito.petsc.iet.utils import petsc_call, petsc_call_mpi
 
 
@@ -26,7 +35,6 @@ def lower_petsc(iet, **kwargs):
         return iet, {}
 
     metadata = core_metadata()
-
     data = FindNodes(PetscMetaData).visit(iet)
 
     if any(filter(lambda i: isinstance(i.expr.rhs, Initialize), data)):
@@ -35,10 +43,10 @@ def lower_petsc(iet, **kwargs):
     if any(filter(lambda i: isinstance(i.expr.rhs, Finalize), data)):
         return finalize(iet), metadata
 
-    targets = [i.expr.rhs.target for (i,) in injectsolve_mapper.values()]
-
-    # Assumption is that all targets have the same grid so can use any target here
-    objs = build_core_objects(targets[-1], **kwargs)
+    unique_grids = {i.expr.rhs.grid for (i,) in injectsolve_mapper.values()}
+    # Assumption is that all solves are on the same grid
+    if len(unique_grids) > 1:
+        raise ValueError("All PETScSolves must use the same Grid, but multiple found.")
 
     # Create core PETSc calls (not specific to each PETScSolve)
     core = make_core_petsc_calls(objs, **kwargs)
@@ -54,17 +62,18 @@ def lower_petsc(iet, **kwargs):
         setup.extend(builder.solversetup.calls)
 
         # Transform the spatial iteration loop with the calls to execute the solver
-        subs.update(builder.solve.mapper)
+        subs.update({builder.solve.spatial_body: builder.solve.calls})
 
         efuncs.update(builder.cbbuilder.efuncs)
 
+    populate_matrix_context(efuncs, objs)
+
     iet = Transformer(subs).visit(iet)
 
     body = core + tuple(setup) + (BlankLine,) + iet.body.body
     body = iet.body._rebuild(body=body)
     iet = iet._rebuild(body=body)
     metadata.update({'efuncs': tuple(efuncs.values())})
-
     return iet, metadata
 
 
@@ -100,56 +109,140 @@ def make_core_petsc_calls(objs, **kwargs):
     return call_mpi, BlankLine
 
 
-def build_core_objects(target, **kwargs):
-    communicator = 'PETSC_COMM_WORLD'
-
-    return {
-        'size': PetscMPIInt(name='size'),
-        'comm': communicator,
-        'err': PetscErrorCode(name='err'),
-        'grid': target.grid
-    }
-
-
 class Builder:
     """
     This class is designed to support future extensions, enabling
     different combinations of solver types, preconditioning methods,
     and other functionalities as needed.
-
     The class will be extended to accommodate different solver types by
     returning subclasses of the objects initialised in __init__,
     depending on the properties of `injectsolve`.
     """
     def __init__(self, injectsolve, objs, iters, **kwargs):
+        self.injectsolve = injectsolve
+        self.objs = objs
+        self.iters = iters
+        self.kwargs = kwargs
+        self.coupled = isinstance(injectsolve.expr.rhs.fielddata, MultipleFieldData)
+        self.args = {
+            'injectsolve': self.injectsolve,
+            'objs': self.objs,
+            'iters': self.iters,
+            **self.kwargs
+        }
+        self.args['solver_objs'] = self.objbuilder.solver_objs
+        self.args['timedep'] = self.timedep
+        self.args['cbbuilder'] = self.cbbuilder
+
+    @cached_property
+    def objbuilder(self):
+        return (
+            CoupledObjectBuilder(**self.args)
+            if self.coupled else
+            BaseObjectBuilder(**self.args)
+        )
 
-        # Determine the time dependency class
-        time_mapper = injectsolve.expr.rhs.time_mapper
-        timedep = TimeDependent if time_mapper else NonTimeDependent
-        self.timedep = timedep(injectsolve, iters, **kwargs)
+    @cached_property
+    def timedep(self):
+        time_mapper = self.injectsolve.expr.rhs.time_mapper
+        timedep_class = TimeDependent if time_mapper else NonTimeDependent
+        return timedep_class(**self.args)
 
-        # Objects
-        self.objbuilder = BaseObjectBuilder(injectsolve, **kwargs)
-        self.solver_objs = self.objbuilder.solver_objs
+    @cached_property
+    def cbbuilder(self):
+        return CCBBuilder(**self.args) if self.coupled else CBBuilder(**self.args)
 
-        # Callbacks
-        self.cbbuilder = CallbackBuilder(
-            injectsolve, objs, self.solver_objs, timedep=self.timedep,
-            **kwargs
-        )
+    @cached_property
+    def solversetup(self):
+        return CoupledSetup(**self.args) if self.coupled else BaseSetup(**self.args)
 
-        # Solver setup
-        self.solversetup = BaseSetup(
-            self.solver_objs, objs, injectsolve, self.cbbuilder
-        )
+    @cached_property
+    def solve(self):
+        return CoupledSolver(**self.args) if self.coupled else Solver(**self.args)
 
-        # Execute the solver
-        self.solve = Solver(
-            self.solver_objs, objs, injectsolve, iters,
-            self.cbbuilder, timedep=self.timedep
-        )
+
+def populate_matrix_context(efuncs, objs):
+    if not objs['dummyefunc'] in efuncs.values():
+        return
+
+    subdms_expr = DummyExpr(
+        FieldFromPointer(objs['Subdms']._C_symbol, objs['ljacctx']),
+        objs['Subdms']._C_symbol
+    )
+    fields_expr = DummyExpr(
+        FieldFromPointer(objs['Fields']._C_symbol, objs['ljacctx']),
+        objs['Fields']._C_symbol
+    )
+    body = CallableBody(
+        List(body=[subdms_expr, fields_expr]),
+        init=(objs['begin_user'],),
+        retstmt=tuple([Call('PetscFunctionReturn', arguments=[0])])
+    )
+    name = 'PopulateMatContext'
+    efuncs[name] = Callable(
+        name, body, objs['err'],
+        parameters=[objs['ljacctx'], objs['Subdms'], objs['Fields']]
+    )
 
 
-# Move these to types folder
+# TODO: Devito MPI + PETSc testing
+# if kwargs['options']['mpi'] -> communicator = grid.distributor._obj_comm
+communicator = 'PETSC_COMM_WORLD'
+subdms = PointerDM(name='subdms')
+fields = PointerIS(name='fields')
+submats = PointerMat(name='submats')
+rows = PointerIS(name='rows')
+cols = PointerIS(name='cols')
+
+
+# A static dict containing shared symbols and objects that are not
+# unique to each PETScSolve.
+# Many of these objects are used as arguments in callback functions to make
+# the C code cleaner and more modular. This is also a step toward leveraging
+# Devito's `reuse_efuncs` functionality, allowing reuse of efuncs when
+# they are semantically identical.
+objs = frozendict({
+    'size': PetscMPIInt(name='size'),
+    'comm': communicator,
+    'err': PetscErrorCode(name='err'),
+    'block': CallbackMat('block'),
+    'submat_arr': PointerMat(name='submat_arr'),
+    'subblockrows': PetscInt('subblockrows'),
+    'subblockcols': PetscInt('subblockcols'),
+    'rowidx': PetscInt('rowidx'),
+    'colidx': PetscInt('colidx'),
+    'J': Mat('J'),
+    'X': GlobalVec('X'),
+    'xloc': LocalVec('xloc'),
+    'Y': GlobalVec('Y'),
+    'yloc': LocalVec('yloc'),
+    'F': GlobalVec('F'),
+    'floc': LocalVec('floc'),
+    'B': GlobalVec('B'),
+    'nfields': PetscInt('nfields'),
+    'irow': PointerIS(name='irow'),
+    'icol': PointerIS(name='icol'),
+    'nsubmats': Scalar('nsubmats', dtype=np.int32),
+    'matreuse': MatReuse('scall'),
+    'snes': SNES('snes'),
+    'rows': rows,
+    'cols': cols,
+    'Subdms': subdms,
+    'LocalSubdms': CallbackPointerDM(name='subdms'),
+    'Fields': fields,
+    'LocalFields': CallbackPointerIS(name='fields'),
+    'Submats': submats,
+    'ljacctx': JacobianStruct(
+        fields=[subdms, fields, submats], modifier=' *'
+    ),
+    'subctx': SubMatrixStruct(fields=[rows, cols]),
+    'Null': Macro('NULL'),
+    'dummyctx': Symbol('lctx'),
+    'dummyptr': DummyArg('dummy'),
+    'dummyefunc': Symbol('dummyefunc'),
+    'dof': PetscInt('dof'),
+    'begin_user': c.Line('PetscFunctionBeginUser;'),
+})
+
+# Move to macros file?
 Null = Macro('NULL')
-void = 'void'