compiler: Add new petsc types for utilising petscsections

Author: ZoeLeibowitz

devito/petsc/iet/builder.py

@@ -345,6 +345,7 @@ class ConstrainedBCMixin:
     """
     """
     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?
@@ -355,10 +356,34 @@ def _create_dmda_calls(self, dmda):
         dm_mat_type = petsc_call('DMSetMatType', [dmda, 'MATSHELL'])
 
         set_constraints = petsc_call(
-            self.callback_builder._constrain_bc_efunc.name, []
+            self.callback_builder._constrain_bc_efunc.name, [dmda]
+        )
+
+        get_local_section = petsc_call('DMGetLocalSection', [dmda, Byref(sobjs['lsection'])])
+
+        get_point_sf = petsc_call('DMGetPointSF', [dmda, Byref(sobjs['sf'])])
+
+        create_global_section = petsc_call(
+            'PetscSectionCreateGlobalSection', [sobjs['lsection'], sobjs['sf'], 'PETSC_TRUE', 'PETSC_FALSE', 'PETSC_FALSE', Byref(sobjs['gsection'])]
+        )
+
+        dm_set_global_section = petsc_call('DMSetGlobalSection', [dmda, sobjs['gsection']])
+
+        dm_create_section_sf = petsc_call('DMCreateSectionSF', [dmda, sobjs['lsection'], sobjs['gsection']])
+
+        return (
+            dmda_create,
+            da_create_section,
+            dm_set_from_opts,
+            dm_setup,
+            dm_mat_type,
+            set_constraints,
+            get_local_section,
+            get_point_sf,
+            create_global_section,
+            dm_set_global_section,
+            dm_create_section_sf
         )
-        # OBVS CLEANUP
-        return dmda_create, da_create_section, dm_set_from_opts, dm_setup, dm_mat_type, set_constraints
 
 
 class ConstrainedBCBuilder(ConstrainedBCMixin, BuilderBase):

devito/petsc/iet/passes.py

@@ -22,7 +22,7 @@
     BaseCallbackBuilder, CoupledCallbackBuilder, populate_matrix_context,
     get_user_struct_fields
 )
-from devito.petsc.iet.type_builder import BaseTypeBuilder, CoupledTypeBuilder, objs
+from devito.petsc.iet.type_builder import BaseTypeBuilder, CoupledTypeBuilder, ConstrainedBCTypeBuilder, objs
 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
@@ -264,11 +264,14 @@ def __init__(self, inject_solve, iters, comm, section_mapper, **kwargs):
 
     @cached_property
     def type_builder(self):
-        return (
-            CoupledTypeBuilder(**self.common_kwargs)
-            if self.coupled else
-            BaseTypeBuilder(**self.common_kwargs)
-        )
+        if self.coupled and self.constrain_bc:
+            return NotImplementedError
+        elif self.coupled:
+            return CoupledTypeBuilder(**self.common_kwargs) 
+        elif self.constrain_bc:
+            return ConstrainedBCTypeBuilder(**self.common_kwargs) 
+        else:
+            return BaseTypeBuilder(**self.common_kwargs) 
 
     @cached_property
     def time_dependence(self):
@@ -280,17 +283,12 @@ def time_dependence(self):
     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):
         if self.coupled and self.constrain_bc:
-            # TODO: implement CoupledConstrainedBCBuilder
-            return CoupledBuilder(**self.common_kwargs) 
+            # TODO: implement CoupledConstrainedBCBuilder 
+            return NotImplementedError
         elif self.coupled:
             return CoupledBuilder(**self.common_kwargs) 
         elif self.constrain_bc:

devito/petsc/iet/type_builder.py

@@ -8,7 +8,7 @@
     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
+    CallbackPointerIS, CallbackMat, DummyArg, NofSubMats, PetscSectionGlobal, PetscSectionLocal, PetscSF
 )
 
 
@@ -199,6 +199,21 @@ def _target_dependent(self, base_dict):
             )
 
 
+class ConstrainedBCTypeBuilder(BaseTypeBuilder):
+    def _extend_build(self, base_dict):
+        sreg = self.sregistry
+        base_dict['lsection'] = PetscSectionLocal(
+            name=sreg.make_name(prefix='lsection')
+        )
+        base_dict['gsection'] = PetscSectionGlobal(
+            name=sreg.make_name(prefix='gsection')
+        )
+        base_dict['sf'] = PetscSF(
+            name=sreg.make_name(prefix='sf')
+        )
+        return base_dict
+
+
 subdms = PointerDM(name='subdms')
 fields = PointerIS(name='fields')
 submats = PointerMat(name='submats')

devito/petsc/types/dimension.py

@@ -29,7 +29,7 @@ def _arg_values(self, grid=None, **kwargs):
         g_x_M = decomp.glb_max
         # print(g_x_M)
         val = decomp.index_glb_to_loc_unsafe(g_x_M - grtkn)
-        print(val)
+        # print(val)
 
 
         return {self.name: int(val)}

devito/petsc/types/object.py

@@ -193,6 +193,22 @@ class SingleIS(PetscObject):
     dtype = CustomDtype('IS')
 
 
+class PetscSectionGlobal(PetscObject):
+    dtype = CustomDtype('PetscSection')
+
+    @property
+    def _C_free(self):
+        return petsc_call('PetscSectionDestroy', [Byref(self.function)])
+
+
+class PetscSectionLocal(PetscObject):
+    dtype = CustomDtype('PetscSection')
+
+
+class PetscSF(PetscObject):
+    dtype = CustomDtype('PetscSF')
+
+
 class PETScStruct(LocalCompositeObject):
 
     @property
@@ -340,5 +356,6 @@ class NofSubMats(Scalar, LocalType):
     Vec: 1,
     Mat: 2,
     SNES: 3,
-    DM: 4,
+    PetscSectionGlobal: 4,
+    DM: 5,
 }