compiler: Some development with the bc callback automation

Author: ZoeLeibowitz

devito/petsc/equations.py

@@ -1,6 +1,6 @@
 from sympy import Eq
 from devito.symbolics import retrieve_indexed, retrieve_dimensions
-from devito.petsc.types.equation import ConstrainEssentialBC
+from devito.petsc.types.equation import ConstrainBC
 from devito.types.dimension import CustomBoundSubDimension
 from devito import Min, Max
 
@@ -22,7 +22,8 @@ def constrain_essential_bcs(expressions, **kwargs):
 
     # build mapper
     for e in expressions:
-        if not isinstance(e, ConstrainEssentialBC):
+        # from IPython import embed; embed()
+        if not isinstance(e, ConstrainBC):
             continue
 
         indexeds = retrieve_indexed(e)
@@ -64,7 +65,7 @@ def constrain_essential_bcs(expressions, **kwargs):
 
     # build new expressions
     for e in expressions:
-        if isinstance(e, ConstrainEssentialBC):
+        if isinstance(e, ConstrainBC):
             new_e = e.subs(mapper)
             new_exprs.append(new_e)
 

devito/petsc/iet/builder.py

@@ -3,7 +3,7 @@
 
 from devito.ir.iet import DummyExpr, BlankLine
 from devito.symbolics import (Byref, FieldFromPointer, VOID,
-                              FieldFromComposite, Null)
+                              FieldFromComposite, Null, String)
 
 from devito.petsc.iet.nodes import (
     FormFunctionCallback, MatShellSetOp, PETScCall, petsc_call
@@ -343,20 +343,25 @@ def _setup(self):
 
 class ConstrainedBCMixin:
     """
+    not really a mixin?
     """
     def _create_dmda_calls(self, dmda):
         sobjs = self.solver_objs
         # 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])
+        da_create_section = petsc_call('PetscOptionsSetValue', [Null, String("-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, [dmda]
+        count_bcs = petsc_call(
+            self.callback_builder._count_bc_efunc.name, [dmda, Byref(sobjs['numBC'])]
+        )
+
+        set_point_bcs = petsc_call(
+            self.callback_builder._point_bc_efunc.name, [dmda, Byref(sobjs['numBC'])]
         )
 
         get_local_section = petsc_call('DMGetLocalSection', [dmda, Byref(sobjs['lsection'])])
@@ -377,7 +382,8 @@ def _create_dmda_calls(self, dmda):
             dm_set_from_opts,
             dm_setup,
             dm_mat_type,
-            set_constraints,
+            count_bcs,
+            set_point_bcs,
             get_local_section,
             get_point_sf,
             create_global_section,

devito/petsc/iet/callbacks.py

@@ -17,6 +17,7 @@
 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
+from devito.petsc.types.object import TempSymb
 
 
 class BaseCallbackBuilder:
@@ -44,7 +45,8 @@ def __init__(self, **kwargs):
         self._user_struct_callback = None
         self._F_efunc = None
         self._b_efunc = None
-        self._constrain_bc_efunc = None
+        self._count_bc_efunc = None
+        self._point_bc_efunc = None
 
         self._J_efuncs = []
         # TODO: isn't there only ever one of these per solver so why is it a list?
@@ -649,25 +651,39 @@ def _create_initial_guess_body(self, body):
         return Uxreplace(subs).visit(body)
 
     def _make_constrain_bc(self):
-        exprs = self.field_data.constrain_bc.exprs
+        increment_exprs = self.field_data.constrain_bc.increment_exprs
+        point_bc_exprs = self.field_data.constrain_bc.point_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,
+        irs0, _ = self.rcompile(
+            increment_exprs, options={'mpi': False}, sregistry=self.sregistry,
             concretize_mapper=self.concretize_mapper
         )
-        body = self._create_constrain_bc_body(
-            List(body=irs.uiet.body)
+        # Compile constrain `eqns` into an IET via recursive compilation
+        irs1, _ = self.rcompile(
+            point_bc_exprs, options={'mpi': False}, sregistry=self.sregistry,
+            concretize_mapper=self.concretize_mapper
         )
-        cb = self._make_petsc_callable(
-            'ConstrainBCs', body, parameters=(sobjs['callbackdm'],)
+        count_bc_body = self._create_count_bc_body(
+            List(body=irs0.uiet.body)
         )
-        self._constrain_bc_efunc = cb
-        self._efuncs[cb.name] = cb
+        set_point_bc_body = self._create_set_point_bc_body(
+            List(body=irs1.uiet.body)
+        )
+        cb0 = self._make_petsc_callable(
+            'CountBCs', count_bc_body, parameters=(sobjs['callbackdm'], sobjs['numBCPtr'])
+        )
+        cb1 = self._make_petsc_callable(
+            'SetPointBCs', set_point_bc_body, parameters=(sobjs['callbackdm'], sobjs['numBC'])
+        )
+        self._count_bc_efunc = cb0
+        self._efuncs[cb0.name] = cb0
+        self._point_bc_efunc = cb1
+        self._efuncs[cb1.name] = cb1
 
-    def _create_constrain_bc_body(self, body):
+    def _create_count_bc_body(self, body):
         linsolve_expr = self.inject_solve.expr.rhs
         objs = self.objs
         sobjs = self.solver_objs
@@ -676,11 +692,6 @@ def _create_constrain_bc_body(self, body):
         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)]
@@ -695,15 +706,14 @@ def _create_constrain_bc_body(self, body):
             'DMGetApplicationContext', [dmda, Byref(ctx._C_symbol)]
         )
 
-        vec_restore_array = petsc_call(
-            'VecRestoreArray', [objs['xloc'], Byref(x_arr._C_symbol)]
-        )
+        # dummyexpr = Dereference(self.target, sobjs['numBCPtr'])
 
-        body = body._rebuild(body=body.body + (vec_restore_array,))
+        # body = body._rebuild(body=body.body)
+
+        body = body._rebuild(body.body)
 
         stacks = (
-            vec_get_array,
-            dm_get_local_info
+            dm_get_local_info,
         )
 
         # Dereference function data in struct
@@ -720,6 +730,62 @@ def _create_constrain_bc_body(self, body):
         # Replace non-function data with pointer to data in struct
         subs = {i._C_symbol: FieldFromPointer(i._C_symbol, ctx) for
                 i in fields if not isinstance(i.function, AbstractFunction)}
+        
+        # subs[]
+        # subs[self.target] = sobjs['numBC']
+        
+        subs[TempSymb._C_symbol] = sobjs['numBCPtr']._C_symbol
+
+        # from IPython import embed; embed()
+
+        return Uxreplace(subs).visit(body)
+    
+    def _create_set_point_bc_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_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)]
+        )
+
+        # body = body._rebuild(body=body.body)
+
+        stacks = (
+            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, Definition(sobjs['k_iter'])]
+
+        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)}
+        
+
+        subs[TempSymb._C_symbol] = sobjs['bcPointsArr'].indexed[sobjs['k_iter']]
 
         return Uxreplace(subs).visit(body)
 

devito/petsc/iet/type_builder.py

@@ -8,7 +8,8 @@
     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, PetscSectionGlobal, PetscSectionLocal, PetscSF
+    CallbackPointerIS, CallbackMat, DummyArg, NofSubMats, PetscSectionGlobal, PetscSectionLocal, PetscSF,
+    PetscIntPtr, CallbackPetscInt, CallbackPointerPetscInt
 )
 
 
@@ -210,6 +211,19 @@ def _extend_build(self, base_dict):
         )
         base_dict['sf'] = PetscSF(
             name=sreg.make_name(prefix='sf')
+        )  
+        name = sreg.make_name(prefix='numBC')
+        base_dict['numBC'] = PetscInt(
+            name=name, initvalue=0
+        )
+        base_dict['numBCPtr'] = CallbackPetscInt(
+            name=sreg.make_name(prefix='numBCPtr'), initvalue=0
+        )
+        base_dict['bcPointsArr'] = CallbackPointerPetscInt(
+            name=sreg.make_name(prefix='bcPointsArr')
+        )
+        base_dict['k_iter'] = PetscInt(
+            name='k_iter', initvalue=0
         )
         return base_dict
 

devito/petsc/types/equation.py

@@ -1,4 +1,4 @@
-from devito.types.equation import Eq
+from devito.types.equation import Eq, Inc
 
 
 __all__ = ['EssentialBC']
@@ -45,10 +45,28 @@ class ZeroColumn(EssentialBC):
     pass
 
 
-class ConstrainEssentialBC(EssentialBC):
-    """
-    Equation used to constrain nodes marked by EssentialBCs
-    inside a PetscSection. This type of equation is generated inside
-    petscsolve if the user sets `constrain_bcs=True`.
-    """
+class ConstrainBC(EssentialBC):
     pass
+
+
+# class NoOfEssentialBC(ConstrainBC, Inc):
+#     """
+#     Equation used count essential boundary condition nodes.
+#     This type of equation is generated inside
+#     petscsolve if the user sets `constrain_bcs=True`.
+#     """
+#     def __new__(cls, *args, **kwargs):
+#         return Inc.__new__(Inc, *args, **kwargs)
+    
+
+class NoOfEssentialBC(ConstrainBC, Inc):
+    """Equation used count essential boundary condition nodes.
+    This type of equation is generated inside
+    petscsolve if the user sets `constrain_bcs=True`."""
+    def __new__(cls, *args, **kwargs):
+        return Inc.__new__(Inc, *args, **kwargs)
+
+
+class PointEssentialBC(ConstrainBC):
+    pass
+