compiler: Progress with petscsection constrain bc callback

Author: ZoeLeibowitz

devito/operator/operator.py

@@ -371,6 +371,7 @@ def _lower_exprs(cls, expressions, **kwargs):
 
         # rename etc
         expressions = lower_exprs_petsc(expressions, **kwargs)
+        # from IPython import embed; embed()
 
         processed = [LoweredEq(i) for i in expressions]
 

devito/petsc/equations.py

@@ -1,19 +1,30 @@
 from sympy import Eq
 from devito.symbolics import retrieve_indexed, retrieve_dimensions
-from devito.petsc import EssentialBC
+from devito.petsc.types.equation import ConstrainEssentialBC
 from devito.types.dimension import CustomBoundSubDimension
 from devito import Min, Max
 
 
 def lower_exprs_petsc(expressions, **kwargs):
-    mapper = {}
+    # Constrain EssentialBCs using PetscSection if specified to do so
+    expressions = constrain_essential_bcs(expressions, **kwargs)
+
+    return expressions
+
 
-    additional_exprs = []
+
+def constrain_essential_bcs(expressions, **kwargs):
+    """TODO: improve docs ..Modify the subdims used in ConstrainEssentialBC equations ... to locally
+    constrain nodes (including non owned halo nodes) ....."""
+
+    mapper = {}
+    new_exprs = []
 
     # build mapper
     for e in expressions:
-        if not isinstance(e, EssentialBC):
+        if not isinstance(e, ConstrainEssentialBC):
             continue
+
         indexeds = retrieve_indexed(e)
         dims = retrieve_dimensions([i for j in indexeds for i in j.indices], mode='unique')
 
@@ -30,8 +41,6 @@ def lower_exprs_petsc(expressions, **kwargs):
 
             from devito.petsc.types.dimension import SubDimMax, SubDimMin
 
-
-
             # TODO: change name..
 
             # in theory this class shoulod just take in d
@@ -51,20 +60,13 @@ def lower_exprs_petsc(expressions, **kwargs):
             )
             mapper[d] = new_dim
 
-            # from IPython import embed; embed()
-
     # build new expressions
     for e in expressions:
-        if not isinstance(e, EssentialBC):
-            continue
-
-        # build new expression
-        new_e = e.subs(mapper)
-
-        additional_exprs.append(new_e)
-
-    # return expressions + additional_exprs
-    return expressions + additional_exprs
-
-
-
+        if isinstance(e, ConstrainEssentialBC):
+            new_e = e.subs(mapper)
+            new_exprs.append(new_e)
+        
+        else:
+            new_exprs.append(e)
+
+    return new_exprs

devito/petsc/iet/builder.py

@@ -1,4 +1,5 @@
 import math
+from functools import cached_property
 
 from devito.ir.iet import DummyExpr, BlankLine
 from devito.symbolics import (Byref, FieldFromPointer, VOID,
@@ -22,7 +23,11 @@ def __init__(self, **kwargs):
         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()
+        # self.calls = self._setup()
+
+    @cached_property
+    def calls(self):
+        return self._setup()
 
     @property
     def snes_ctx(self):
@@ -142,9 +147,11 @@ def _extend_setup(self):
 
     def _create_dmda_calls(self, dmda):
         dmda_create = self._create_dmda(dmda)
+        # TODO: probs need to set the dm options prefix the same as snes?
+        dm_set_from_opts = petsc_call('DMSetFromOptions', [dmda])
         dm_setup = petsc_call('DMSetUp', [dmda])
         dm_mat_type = petsc_call('DMSetMatType', [dmda, 'MATSHELL'])
-        return dmda_create, dm_setup, dm_mat_type
+        return dmda_create, dm_set_from_opts, dm_setup, dm_mat_type
 
     def _create_dmda(self, dmda):
         sobjs = self.solver_objs
@@ -184,7 +191,7 @@ def _create_dmda(self, dmda):
         dmda = petsc_call(f'DMDACreate{nspace_dims}d', args)
 
         return dmda
-
+    
 
 class CoupledBuilder(BuilderBase):
     def _setup(self):
@@ -332,6 +339,35 @@ def _setup(self):
             create_submats) + \
             tuple(deref_dms) + tuple(xglobals) + tuple(xlocals) + (BlankLine,)
         return coupled_setup
+    
+
+class ConstrainedBCMixin:
+    """
+    """
+    def _create_dmda_calls(self, dmda):
+        # TODO: CLEAN UP
+        dmda_create = self._create_dmda(dmda)
+        # TODO: probs need to set the dm options prefix the same as snes?
+        # don't hardcode this probs? - the dm needs to be specific to the solver as well
+        da_create_section = petsc_call('PetscOptionsSetValue', [Null, '-da_use_section', Null])
+        dm_set_from_opts = petsc_call('DMSetFromOptions', [dmda])
+        dm_setup = petsc_call('DMSetUp', [dmda])
+        dm_mat_type = petsc_call('DMSetMatType', [dmda, 'MATSHELL'])
+
+        set_constraints = petsc_call(
+            self.callback_builder._constrain_bc_efunc.name, []
+        )
+        # OBVS CLEANUP
+        return dmda_create, da_create_section, dm_set_from_opts, dm_setup, dm_mat_type, set_constraints
+    
+
+class ConstrainedBCBuilder(ConstrainedBCMixin, BuilderBase):
+    pass
+
+
+# TODO: Implement this properly
+class CoupledConstrainedBCBuilder(ConstrainedBCMixin, CoupledBuilder):
+    pass
 
 
 def petsc_call_mpi(specific_call, call_args):

devito/petsc/iet/callbacks.py

@@ -37,15 +37,18 @@ def __init__(self, **kwargs):
         self._efuncs = OrderedDict()
         self._struct_params = []
 
+        # TODO: use either efunc or callback lingo here
         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._constrain_bc_efunc = None
 
         self._J_efuncs = []
-        self._initial_guesses = []
+        # TODO: isn't there only ever one of these per solver so why is it a list?
+        self._initial_guess_efuncs = []
 
         self._make_core()
         self._efuncs = self._uxreplace_efuncs()
@@ -82,9 +85,10 @@ def J_efuncs(self):
         """
         return self._J_efuncs
 
+    # TODO: do i really need a property for this - probs not?
     @property
-    def initial_guesses(self):
-        return self._initial_guesses
+    def initial_guess_efuncs(self):
+        return self._initial_guess_efuncs
 
     @property
     def user_struct_callback(self):
@@ -112,11 +116,18 @@ def target(self):
 
     def _make_core(self):
         self._make_options_callback()
+        # Make the mat-vec callback to form the matfree Jacobian
         self._make_matvec(self.field_data.jacobian)
+        # Make the residual callback
         self._make_formfunc()
+        # Make the RHS callback
         self._make_formrhs()
+        # Make the initial guess callback
         if self.field_data.initial_guess.exprs:
             self._make_initial_guess()
+        # Make the callback used to constrain boundary nodes
+        if self.field_data.constrain_bc.exprs:
+            self._make_constrain_bc()
         self._make_user_struct_callback()
 
     def _make_petsc_callable(self, prefix, body, parameters=()):
@@ -578,7 +589,7 @@ def _make_initial_guess(self):
         cb = self._make_petsc_callable(
             'FormInitialGuess', body, parameters=(sobjs['callbackdm'], objs['xloc'])
         )
-        self._initial_guesses.append(cb)
+        self._initial_guess_efuncs.append(cb)
         self._efuncs[cb.name] = cb
 
     def _create_initial_guess_body(self, body):
@@ -636,6 +647,29 @@ def _create_initial_guess_body(self, body):
                 i in fields if not isinstance(i.function, AbstractFunction)}
 
         return Uxreplace(subs).visit(body)
+    
+    def _make_constrain_bc(self):
+        exprs = self.field_data.constrain_bc.exprs
+        sobjs = self.solver_objs
+        objs = self.objs
+
+        # Compile constrain `eqns` into an IET via recursive compilation
+        irs, _ = self.rcompile(
+            exprs, options={'mpi': False}, sregistry=self.sregistry,
+            concretize_mapper=self.concretize_mapper
+        )
+        # from IPython import embed; embed()
+        body = self._create_constrain_bc_body(
+            List(body=irs.uiet.body)
+        )
+        cb = self._make_petsc_callable(
+            'ConstrainBCs', body, parameters=(sobjs['callbackdm'],)
+        )
+        self._constrain_bc_efunc = cb
+        self._efuncs[cb.name] = cb
+
+    def _create_constrain_bc_body(self, body):
+        return body
 
     def _make_user_struct_callback(self):
         """

devito/petsc/iet/passes.py

@@ -23,7 +23,7 @@
     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.builder import BuilderBase, CoupledBuilder, ConstrainedBCBuilder, 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
@@ -248,6 +248,7 @@ def __init__(self, inject_solve, iters, comm, section_mapper, **kwargs):
         self.get_info = inject_solve.expr.rhs.get_info
         self.kwargs = kwargs
         self.coupled = isinstance(inject_solve.expr.rhs.field_data, MultipleFieldData)
+        self.constrain_bc = inject_solve.expr.rhs.field_data.constrain_bc
         self.common_kwargs = {
             'inject_solve': self.inject_solve,
             'objs': self.objs,
@@ -280,10 +281,22 @@ 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 builder(self):
-        return CoupledBuilder(**self.common_kwargs) \
-            if self.coupled else BuilderBase(**self.common_kwargs)
+        if self.coupled and self.constrain_bc:
+            # TODO: implement CoupledConstrainedBCBuilder
+            return CoupledBuilder(**self.common_kwargs) 
+        elif self.coupled:
+            return CoupledBuilder(**self.common_kwargs) 
+        elif self.constrain_bc:
+            return ConstrainedBCBuilder(**self.common_kwargs) 
+        else:
+            return BuilderBase(**self.common_kwargs) 
 
     @cached_property
     def solve(self):

devito/petsc/iet/solve.py

@@ -37,8 +37,8 @@ def _execute_solve(self):
 
         vec_place_array = self.time_dependence.place_array(target)
 
-        if self.callback_builder.initial_guesses:
-            initguess = self.callback_builder.initial_guesses[0]
+        if self.callback_builder.initial_guess_efuncs:
+            initguess = self.callback_builder.initial_guess_efuncs[0]
             initguess_call = petsc_call(initguess.name, [dmda, sobjs['xlocal']])
         else:
             initguess_call = None

devito/petsc/solve.py

@@ -6,7 +6,7 @@
 
 from devito.petsc.types import (
     LinearSolverMetaData, PETScArray, DMDALocalInfo, FieldData, MultipleFieldData,
-    Jacobian, Residual, MixedResidual, MixedJacobian, InitialGuess
+    Jacobian, Residual, MixedResidual, MixedJacobian, InitialGuess, ConstrainBC
 )
 from devito.petsc.types.equation import EssentialBC
 from devito.petsc.solver_parameters import (
@@ -18,7 +18,7 @@
 
 
 def petscsolve(target_exprs, target=None, solver_parameters=None,
-               options_prefix=None, get_info=[]):
+               options_prefix=None, get_info=[], constrain_bcs=False):
     """
     Returns a symbolic expression representing a linear PETSc solver,
     enriched with all the necessary metadata for execution within an `Operator`.
@@ -78,6 +78,12 @@ def petscsolve(target_exprs, target=None, solver_parameters=None,
         - ['kspgetiterationnumber', 'kspgettolerances', 'kspgetconvergedreason',
            'kspgettype', 'kspgetnormtype', 'snesgetiterationnumber']
 
+    constrain_bcs : bool, optional
+        If `True`, essential boundary conditions specifed by `EssentialBC` equations
+        are constrained through a `PetscSection`. As a result, the corresponding degrees
+        of freedom are excluded from the global solver and are not imposed using
+        trivial equations.
+
     Returns
     -------
     Eq:
@@ -86,22 +92,24 @@ def petscsolve(target_exprs, target=None, solver_parameters=None,
     """
     if target is not None:
         return InjectSolve(solver_parameters, {target: target_exprs},
-                           options_prefix, get_info).build_expr()
+                           options_prefix, get_info, constrain_bcs).build_expr()
     else:
+        # TODO: extend mixed case to support constrain_bcs
         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=[]):
+                 get_info=[], constrain_bcs=False):
         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]
+        self.constrain_bcs = constrain_bcs
 
     def build_expr(self):
         target, funcs, field_data = self.linear_solve_args()
@@ -129,16 +137,26 @@ def linear_solve_args(self):
 
         exprs = sorted(exprs, key=lambda e: not isinstance(e, EssentialBC))
 
+        # TODO: rethink about how essential bcs need to be treated if constrain_bcs is enabled
+        # likely don't need various bits of functionality inside these classes if constrain_bcs is enabled
         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)
 
+        # TODO: extend this to mixed case
+        # TODO: clean up
+        if self.constrain_bcs:
+            constrain_bc = ConstrainBC(target, exprs, arrays)
+        else:
+            constrain_bc = None
+
         field_data = FieldData(
             target=target,
             jacobian=jacobian,
             residual=residual,
             initial_guess=initial_guess,
-            arrays=arrays
+            arrays=arrays,
+            constrain_bc=constrain_bc
         )
 
         return target, funcs, field_data