compiler: Support increment indexing

Author: ZoeLeibowitz

devito/ir/cgen/printer.py

@@ -17,7 +17,7 @@
 from devito.arch.compiler import AOMPCompiler
 from devito.symbolics.inspection import has_integer_args, sympy_dtype
 from devito.symbolics.queries import q_leaf
-from devito.types.basic import AbstractFunction
+from devito.types.basic import AbstractFunction, PostIncrementIndex
 from devito.tools import ctypes_to_cstr, dtype_to_ctype, ctypes_vector_mapper
 
 __all__ = ['BasePrinter', 'ccode']
@@ -148,8 +148,13 @@ def _print_Indexed(self, expr):
         --------
         U[t,x,y,z] -> U[t][x][y][z]
         """
-        inds = ''.join(['[' + self._print(x) + ']' for x in expr.indices])
-        return f'{self._print(expr.base.label)}{inds}'
+        inds = []
+        for i in expr.indices:
+            if isinstance(i, PostIncrementIndex):
+                inds.append(f"[{self._print(i)}++]")
+            else:
+                inds.append(f"[{self._print(i)}]")
+        return f"{self._print(expr.base.label)}{''.join(inds)}"
 
     def _print_FIndexed(self, expr):
         """
@@ -165,7 +170,19 @@ def _print_FIndexed(self, expr):
         except AttributeError:
             label = expr.base.label
         return f'{self._print(label)}({inds})'
-
+    
+    # def _print_PostIncrementIndexed(self, expr):
+    #     """
+    #     Print an Indexed as a ...
+
+    #     Examples
+    #     --------
+    #     U[k] -> U[k++]
+    #     """
+    #     # from IPython import embed; embed()
+    #     inds = ''.join(['[' + self._print(x) + '++' + ']' for x in expr.indices])
+    #     return f'{self._print(expr.base.label)}{inds}'
+    
     def _print_Rational(self, expr):
         """Print a Rational as a C-like float/float division."""
         # This method and _print_Float below forcefully add a F to any

devito/petsc/iet/callbacks.py

@@ -5,7 +5,7 @@
     BlankLine, Callable, Iteration, PointerCast, Definition
 )
 from devito.symbolics import (
-    Byref, FieldFromPointer, IntDiv, Deref, Mod, String, Null, VOID
+    Byref, FieldFromPointer, IntDiv, Deref, Mod, String, Null, VOID, Cast
 )
 from devito.symbolics.unevaluation import Mul
 from devito.types.basic import AbstractFunction
@@ -17,7 +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
+from devito.petsc.types.object import Counter
 
 
 class BaseCallbackBuilder:
@@ -693,10 +693,6 @@ def _create_count_bc_body(self, body):
         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)
@@ -710,21 +706,24 @@ def _create_count_bc_body(self, body):
 
         # body = body._rebuild(body=body.body)
 
-        body = body._rebuild(body.body)
-
-        stacks = (
-            dm_get_local_info,
-        )
-
         # Dereference function data in struct
-        derefs = dereference_funcs(ctx, fields)
+        # derefs = dereference_funcs(ctx, fields)
+        
+        # OBVS change names
+        deref_ptr = DummyExpr(Counter, Deref(sobjs['numBCPtr']))
+        move_ptr = DummyExpr(Deref(sobjs['numBCPtr']), Counter)
+
+        # from IPython import embed; embed()
 
         # 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]
 
+
+        body = body._rebuild(body.body + (move_ptr,))
+
         body = self._make_callable_body(
-            body, standalones=struct_definition, stacks=stacks+derefs
+            body, standalones=struct_definition, stacks=(deref_ptr,)
         )
 
         # Replace non-function data with pointer to data in struct
@@ -734,7 +733,7 @@ def _create_count_bc_body(self, body):
         # subs[]
         # subs[self.target] = sobjs['numBC']
 
-        subs[TempSymb._C_symbol] = sobjs['numBCPtr']._C_symbol
+        # subs[Counter._C_symbol] = Cast(Deref(sobjs['numBCPtr']._C_symbol))
 
         # from IPython import embed; embed()
 
@@ -785,7 +784,7 @@ def _create_set_point_bc_body(self, body):
                 i in fields if not isinstance(i.function, AbstractFunction)}
 
 
-        subs[TempSymb._C_symbol] = sobjs['bcPointsArr'].indexed[sobjs['k_iter']]
+        subs[Counter._C_symbol] = sobjs['bcPointsArr'].indexed[sobjs['k_iter']]
 
         return Uxreplace(subs).visit(body)
 

devito/petsc/iet/type_builder.py

@@ -9,7 +9,7 @@
     PointerIS, PointerDM, VecScatter, JacobianStruct, SubMatrixStruct, CallbackDM,
     PetscMPIInt, PetscErrorCode, PointerMat, MatReuse, CallbackPointerDM,
     CallbackPointerIS, CallbackMat, DummyArg, NofSubMats, PetscSectionGlobal, PetscSectionLocal, PetscSF,
-    PetscIntPtr, CallbackPetscInt, CallbackPointerPetscInt
+    PetscIntPtr, CallbackPetscInt, CallbackPointerPetscInt, PostIncrementIndex
 )
 
 
@@ -222,7 +222,7 @@ def _extend_build(self, base_dict):
         base_dict['bcPointsArr'] = CallbackPointerPetscInt(
             name=sreg.make_name(prefix='bcPointsArr')
         )
-        base_dict['k_iter'] = PetscInt(
+        base_dict['k_iter'] = PostIncrementIndex(
             name='k_iter', initvalue=0
         )
         return base_dict

devito/petsc/types/metadata.py

@@ -10,7 +10,7 @@
 from devito.operations.solve import eval_time_derivatives
 
 from devito.petsc.config import petsc_variables
-from devito.petsc.types.object import PetscInt, TempSymb
+from devito.petsc.types.object import PetscInt, Counter
 from devito.petsc.types.equation import (
     EssentialBC, ZeroRow, ZeroColumn, NoOfEssentialBC, PointEssentialBC
 )
@@ -762,12 +762,8 @@ def _make_increment_expr(self, expr):
         """
         if isinstance(expr, EssentialBC):
             assert expr.lhs == self.target
-            # return NoOfEssentialBC(
-            #     TempSymb, expr.rhs,
-            #     subdomain=expr.subdomain,
-            # )
             return NoOfEssentialBC(
-                TempSymb, 1,
+                Counter, 1,
                 subdomain=expr.subdomain,
                 implicit_dims=expr.subdomain.dimensions
             )
@@ -790,11 +786,11 @@ def _make_point_bc_expr(self, expr):
         Make the Eq that is used to increment the number of essential
         boundary nodes in the generated ccode.
         """
-        numBC = PetscInt(name='numBC2')
+        # numBC = PetscInt(name='numBC2')
         if isinstance(expr, EssentialBC):
             assert expr.lhs == self.target
             return PointEssentialBC(
-                TempSymb, expr.rhs,
+                Counter, expr.rhs,
                 subdomain=expr.subdomain
             )
         else:

devito/petsc/types/object.py

@@ -1,12 +1,13 @@
 from ctypes import POINTER, c_char
+from functools import cached_property
 
 from devito.tools import CustomDtype, dtype_to_ctype, as_tuple, CustomIntType
 from devito.types import (
     LocalObject, LocalCompositeObject, ModuloDimension, TimeDimension, ArrayObject,
     CustomDimension, Scalar
 )
 from devito.symbolics import Byref, cast
-from devito.types.basic import DataSymbol, LocalType
+from devito.types.basic import DataSymbol, LocalType, PostIncrementIndex
 
 from devito.petsc.iet.nodes import petsc_call
 
@@ -319,6 +320,10 @@ def dtype(self):
         return CustomDtype('IS', modifier=' *')
 
 
+class PetscPostIncrementIndex(PostIncrementIndex):
+    pass
+    
+
 class CallbackPointerPetscInt(PETScArrayObject):
     """
     """
@@ -373,7 +378,8 @@ class NofSubMats(Scalar, LocalType):
     pass
 
 
-TempSymb = PetscInt(name='numBC2')
+# Can this be attached to the consrain bc object in metadata maybe? probs shoulnd't be here
+Counter = PetscInt(name='count')
 
 
 FREE_PRIORITY = {

devito/symbolics/extended_sympy.py

@@ -947,6 +947,3 @@ def rfunc(func, item, *args):
     min: Min,
     max: Max,
 }
-
-
-Null = Macro('NULL')

devito/types/basic.py

@@ -1927,3 +1927,7 @@ def _mem_internal_lazy(self):
     to impose pass-by-reference semantics.
     """
     _C_modifier = None
+
+
+class PostIncrementIndex(DataSymbol):
+    pass