compiler: Progress with the SetPointBCs0 callback petsc

Author: ZoeLeibowitz

devito/passes/iet/linearization.py

@@ -25,7 +25,7 @@ def linearize(graph, **kwargs):
 
     mode = options.get('linearize')
     maybe_callback = kwargs.pop('callback', mode)
-
+    # from IPython import embed; embed()
     if not maybe_callback:
         return
     elif callable(maybe_callback):
@@ -212,7 +212,7 @@ def linearize_accesses(iet, key0, tracker=None):
     indexeds = FindSymbols('indexeds').visit(iet)
     needs = filter_ordered(i.function for i in indexeds if key0(i.function))
     needs = sorted(needs, key=lambda f: len(f.dimensions), reverse=True)
-
+    # from IPython import embed; embed()
     # Update unique sizes and strides
     tracker.update(needs)
 
@@ -230,9 +230,10 @@ def linearize_accesses(iet, key0, tracker=None):
             continue
 
         v = generate_linearization(f, i, tracker)
+        # from IPython import embed; embed()
         if v is not None:
             subs[i] = v
-
+    # from IPython import embed; embed()
     iet = Uxreplace(subs).visit(iet)
 
     # 2) What `iet` *offers*

devito/petsc/iet/callbacks.py

@@ -11,6 +11,7 @@
 from devito.types.basic import AbstractFunction
 from devito.types import Dimension, Temp, TempArray
 from devito.tools import filter_ordered
+from devito.passes.iet.linearization import linearize_accesses
 
 from devito.petsc.iet.nodes import PETScCallable, MatShellSetOp, petsc_call
 from devito.petsc.types import DMCast, MainUserStruct, CallbackUserStruct
@@ -28,6 +29,7 @@ def __init__(self, **kwargs):
 
         self.rcompile = kwargs.get('rcompile', None)
         self.sregistry = kwargs.get('sregistry', None)
+        self.options = kwargs.get('options', {})
         self.concretize_mapper = kwargs.get('concretize_mapper', {})
         self.time_dependence = kwargs.get('time_dependence')
         self.objs = kwargs.get('objs')
@@ -753,6 +755,25 @@ 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}
+
+        # from IPython import embed; embed()
+
+        # findexeds = 
         body = self.time_dependence.uxreplace_time(body)
 
         fields = get_user_struct_fields(body)
@@ -762,7 +783,12 @@ def _create_set_point_bc_body(self, body):
             'DMGetApplicationContext', [dmda, Byref(ctx._C_symbol)]
         )
 
-        # body = body._rebuild(body=body.body)
+        comm = sobjs['comm']
+        is_create_general = petsc_call(
+            'ISCreateGeneral', [comm, sobjs['numBC'], sobjs['bcPointsArr'], 'PETSC_OWN_POINTER']
+        )
+
+        body = body._rebuild(body=body.body + (is_create_general,))
 
         stacks = (
             dm_get_local_info,
@@ -786,7 +812,10 @@ def _create_set_point_bc_body(self, body):
 
         subs[Counter._C_symbol] = sobjs['bcPointsArr'].indexed[sobjs['k_iter']]
 
-        return Uxreplace(subs).visit(body)
+        body = Uxreplace(mapper_findexeds).visit(body)
+        body = Uxreplace(subs).visit(body)
+
+        return body
 
     def _make_user_struct_callback(self):
         """

devito/types/misc.py

@@ -149,13 +149,31 @@ def bind(self, pname):
         findexed = self.func(accessor=accessor)
 
         return ((define, expr), findexed)
+    
+
+    @property
+    def linear_index(self):
+        f = self.function
+        strides_map = self.strides_map
+        indices = self.indices
+
+        items = [
+            idx * strides_map[d]
+            for idx, d in zip(indices, f.dimensions[1:])
+        ]
+        items.append(indices[-1])
+        
+        return sympy.Add(*items, evaluate=False)
 
     func = Pickable._rebuild
 
     # Pickling support
     __reduce_ex__ = Pickable.__reduce_ex__
 
 
+# the special postindex type sould live in this file i think
+
+
 class Global(Symbol):
 
     """