compiler: Working petscsection ed bueler example apart from the byref with malloc

Author: ZoeLeibowitz

devito/petsc/iet/builder.py

@@ -40,13 +40,7 @@ def snes_ctx(self):
     def _setup(self):
         sobjs = self.solver_objs
         dmda = sobjs['dmda']
-        mainctx = sobjs['userctx']
-
-        call_struct_callback = petsc_call(
-            self.callback_builder.user_struct_callback.name, [Byref(mainctx)]
-        )
-
-        calls_set_app_ctx = petsc_call('DMSetApplicationContext', [dmda, Byref(mainctx)])
+        # mainctx = sobjs['userctx']
 
         snes_create = petsc_call('SNESCreate', [sobjs['comm'], Byref(sobjs['snes'])])
 
@@ -116,8 +110,8 @@ def _setup(self):
         mat_set_dm = petsc_call('MatSetDM', [sobjs['Jac'], dmda])
 
         base_setup = dmda_calls + (
-            call_struct_callback,
-            calls_set_app_ctx,
+            # call_struct_callback,
+            # calls_set_app_ctx,
             snes_create,
             snes_options_prefix,
             set_options,
@@ -150,7 +144,15 @@ def _create_dmda_calls(self, dmda):
         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_set_from_opts, dm_setup, dm_mat_type
+        mainctx = self.solver_objs['userctx']
+
+        call_struct_callback = petsc_call(
+            self.callback_builder.user_struct_callback.name, [Byref(mainctx)]
+        )
+
+        calls_set_app_ctx = petsc_call('DMSetApplicationContext', [dmda, Byref(mainctx)])
+
+        return dmda_create, dm_set_from_opts, dm_setup, dm_mat_type, call_struct_callback, calls_set_app_ctx
 
     def _create_dmda(self, dmda):
         sobjs = self.solver_objs
@@ -344,6 +346,8 @@ class ConstrainedBCMixin:
     """
     def _create_dmda_calls(self, dmda):
         sobjs = self.solver_objs
+        # mainctx = sobjs['mainctx']
+        mainctx = sobjs['userctx']
         # TODO: CLEAN UP
         dmda_create = self._create_dmda(dmda)
         # TODO: probs need to set the dm options prefix the same as snes?
@@ -373,12 +377,21 @@ def _create_dmda_calls(self, dmda):
 
         dm_create_section_sf = petsc_call('DMCreateSectionSF', [dmda, sobjs['lsection'], sobjs['gsection']])
 
+
+        call_struct_callback = petsc_call(
+            self.callback_builder.user_struct_callback.name, [Byref(mainctx)]
+        )
+
+        calls_set_app_ctx = petsc_call('DMSetApplicationContext', [dmda, Byref(mainctx)])
+
         return (
             dmda_create,
             da_create_section,
             dm_set_from_opts,
             dm_setup,
             dm_mat_type,
+            call_struct_callback,
+            calls_set_app_ctx,
             count_bcs,
             set_point_bcs,
             get_local_section,

devito/petsc/iet/callbacks.py

@@ -755,23 +755,22 @@ def _create_set_point_bc_body(self, body):
             'DMDAGetLocalInfo', [dmda, Byref(linsolve_expr.localinfo)]
         )
 
-        import numpy as np
-        if self.options['index-mode'] == 'int32':
-            dtype = np.int32
-        else:
-            dtype = np.int64
-        from devito.passes.iet.linearization import Tracker
-
-        tracker = Tracker('basic', dtype, self.sregistry)
-
-        key = lambda f: f.name == 'bc'
-        body = linearize_accesses(body, key0=key, tracker=tracker)
-
-        # will only be findexeds 'indexeds'
-        findexeds = FindSymbols('indexeds').visit(body)
-        mapper_findexeds = {i: i.linear_index for i in findexeds}
-
-
+        # import numpy as np
+        # if self.options['index-mode'] == 'int32':
+        #     dtype = np.int32
+        # else:
+        #     dtype = np.int64
+        # from devito.passes.iet.linearization import Tracker
+
+        # tracker = Tracker('basic', dtype, self.sregistry)
+        # # from IPython import embed; embed()
+        # key = lambda f: f.name == 'u'
+        # body = linearize_accesses(body, key0=key, tracker=tracker)
+
+        # # will only be findexeds 'indexeds'
+        # findexeds = FindSymbols('indexeds').visit(body)
+        # mapper_findexeds = {i: i.linear_index for i in findexeds}
+        
         # from IPython import embed; embed()
 
         # findexeds = 
@@ -789,7 +788,11 @@ def _create_set_point_bc_body(self, body):
             'ISCreateGeneral', [comm, sobjs['numBC'], sobjs['bcPointsArr'], 'PETSC_OWN_POINTER', Byref(sobjs['bcPointsIS'])]
         )
 
-        malloc = petsc_call(
+        malloc_bc_points_arr = petsc_call(
+            'PetscMalloc1', [sobjs['numBC'], sobjs['bcPointsArr']]
+        )
+
+        malloc_bc_points = petsc_call(
             'PetscMalloc1', [1, sobjs['bcPoints']]
         )
 
@@ -798,7 +801,7 @@ def _create_set_point_bc_body(self, body):
         set_point_bc = petsc_call(
             'DMDASetPointBC', [dmda, 1, sobjs['bcPoints'], Null]
         )
-        body = body._rebuild(body=body.body + (is_create_general, malloc, dummy_expr, set_point_bc))
+        body = body._rebuild(body=(malloc_bc_points_arr,)+ body.body + (is_create_general, malloc_bc_points, dummy_expr, set_point_bc))
 
         stacks = (
             dm_get_local_info,
@@ -821,7 +824,7 @@ def _create_set_point_bc_body(self, body):
 
         subs[Counter._C_symbol] = sobjs['bcPointsArr'].indexed[sobjs['k_iter']]
 
-        body = Uxreplace(mapper_findexeds).visit(body)
+        # body = Uxreplace(mapper_findexeds).visit(body)
         body = Uxreplace(subs).visit(body)
 
         return body

devito/petsc/iet/passes.py

@@ -9,7 +9,9 @@
 )
 from devito.symbolics import Byref, Macro, Null, FieldFromPointer
 from devito.types.basic import DataSymbol
+from devito.types.misc import FIndexed
 import devito.logger
+from devito.passes.iet.linearization import linearize_accesses
 
 from devito.petsc.types import (
     MultipleFieldData, Initialize, Finalize, ArgvSymbol, MainUserStruct,
@@ -126,6 +128,61 @@ def lower_petsc_symbols(iet, **kwargs):
     # Rebuild `MainUserStruct` and update iet accordingly
     rebuild_parent_user_struct(iet, mapper=callback_struct_mapper)
 
+    # from IPython import embed; embed()
+
+
+    iet = linear_indices(iet, **kwargs)
+
+    ############ tmp
+    # import numpy as np
+    # if kwargs['options']['index-mode'] == 'int32':
+    #     dtype = np.int32
+    # else:
+    #     dtype = np.int64
+    # from devito.passes.iet.linearization import Tracker
+
+    # tracker = Tracker('basic', dtype, kwargs['sregistry'])
+
+    # key = lambda f: f.name == 'bc'
+    # body = linearize_accesses(body, key0=key, tracker=tracker)
+
+    # # will only be findexeds 'indexeds'
+    # findexeds = FindSymbols('indexeds').visit(body)
+    # mapper_findexeds = {i: i.linear_index for i in findexeds}
+
+    # iet = 
+
+
+@iet_pass
+def linear_indices(iet, **kwargs):
+
+    if not iet.name.startswith("SetPointBCs"):
+        return iet, {}
+
+    import numpy as np
+    if kwargs['options']['index-mode'] == 'int32':
+        dtype = np.int32
+    else:
+        dtype = np.int64
+    from devito.passes.iet.linearization import Tracker
+
+    tracker = Tracker('basic', dtype, kwargs['sregistry'])
+    # from IPython import embed; embed()
+    key = lambda f: f.name == 'u'
+    iet = linearize_accesses(iet, key0=key, tracker=tracker)
+    # from IPython import embed; embed()
+    # will only be findexeds 'indexeds'
+    findexeds = [i for i in FindSymbols('indexeds').visit(iet) if isinstance(i, FIndexed)]
+    mapper_findexeds = {i: i.linear_index for i in findexeds}
+
+
+    iet = Uxreplace(mapper_findexeds).visit(iet)
+
+
+    # from IPython import embed; embed()
+
+    return iet, {}
+
 
 @iet_pass
 def rebuild_child_user_struct(iet, mapper, **kwargs):

devito/petsc/types/metadata.py

@@ -789,8 +789,9 @@ def _make_point_bc_expr(self, expr):
         # numBC = PetscInt(name='numBC2')
         if isinstance(expr, EssentialBC):
             assert expr.lhs == self.target
+            # from IPython import embed; embed()
             return PointEssentialBC(
-                Counter, expr.rhs,
+                Counter, self.target,
                 subdomain=expr.subdomain
             )
         else:

devito/types/misc.py

@@ -8,8 +8,8 @@
     # Moved in 1.13
     from sympy.core.basic import ordering_of_classes
 
-from devito.types import Array, CompositeObject, Indexed, Symbol, LocalObject
-from devito.types.basic import IndexedData
+from devito.types import Array, CompositeObject, Indexed, Symbol, LocalObject, ArrayObject
+from devito.types.basic import IndexedData, DataSymbol
 from devito.tools import CustomDtype, Pickable, as_tuple, frozendict
 
 __all__ = ['Timer', 'Pointer', 'VolatileInt', 'FIndexed', 'Wildcard', 'Fence',
@@ -172,8 +172,8 @@ def linear_index(self):
 
 
 # the special postindex type sould live in this file i think
-class PostIncrementIndex(Symbol):
-    pass
+class PostIncrementIndex(LocalObject):
+    dtype = np.int32
 
 
 class Global(Symbol):

examples/petsc/Poisson/ed_bueler_2d.py

@@ -111,11 +111,11 @@ def exact(x, y):
 
 with switchconfig(log_level='DEBUG'):
     op = Operator(petsc, language='petsc')
-    # summary = op.apply()
+    summary = op.apply()
     # print(op.arguments())
 
 
-print(op.ccode)
+# print(op.ccode)
 # iters = summary.petsc[('section0', 'poisson_2d')].KSPGetIterationNumber
 
 u_exact = Function(name='u_exact', grid=grid, space_order=2)