compiler: Change PETScArray type to PetscScalar

.github/workflows/pytest-petsc.yml

@@ -0,0 +1,74 @@
+name: CI-petsc
+
+concurrency:
+  group: ${{ github.workflow }}-${{ github.ref }}
+  cancel-in-progress: true
+
+on:
+  # Trigger the workflow on push or pull request,
+  # but only for the master branch
+  push:
+    branches:
+      - master
+  pull_request:
+    branches:
+      - master
+
+jobs:
+  pytest:
+    name: ${{ matrix.name }}-${{ matrix.set }}
+    runs-on: "${{ matrix.os }}"
+
+    env:
+      DOCKER_BUILDKIT: "1"
+      DEVITO_ARCH: "${{ matrix.arch }}"
+      DEVITO_LANGUAGE: ${{ matrix.language }}
+
+    strategy:
+      # Prevent all build to stop if a single one fails
+      fail-fast: false
+
+      matrix:
+        name: [
+           pytest-docker-py39-gcc-noomp
+        ]
+        include:
+        - name: pytest-docker-py39-gcc-noomp
+          python-version: '3.9'
+          os: ubuntu-latest
+          arch: "gcc"
+          language: "C"
+          sympy: "1.12"
+
+    steps:
+    - name: Checkout devito
+      uses: actions/checkout@v4
+
+    - name: Build docker image
+      run: |
+          docker build . --file docker/Dockerfile.devito --tag devito_img --build-arg base=zoeleibowitz/bases:cpu-${{ matrix.arch }} --build-arg petscinstall=petsc
+
+    - name: Set run prefix
+      run: |
+          echo "RUN_CMD=docker run --rm -t -e CODECOV_TOKEN=${{ secrets.CODECOV_TOKEN }} --name testrun devito_img"  >> $GITHUB_ENV
+      id: set-run
+
+    - name: Set tests
+      run : |
+          echo "TESTS=tests/test_petsc.py" >> $GITHUB_ENV
+      id: set-tests
+
+    - name: Check configuration
+      run: |
+        ${{ env.RUN_CMD }} python3 -c "from devito import configuration; print(''.join(['%s: %s \n' % (k, v) for (k, v) in configuration.items()]))"
+
+    - name: Test with pytest
+      run: |
+        ${{ env.RUN_CMD }} mpiexec -n 1 pytest --cov --cov-config=.coveragerc --cov-report=xml ${{ env.TESTS }}
+
+    - name: Upload coverage to Codecov
+      if: "!contains(matrix.name, 'docker')"
+      uses: codecov/codecov-action@v4
+      with:
+        token: ${{ secrets.CODECOV_TOKEN }}
+        name: ${{ matrix.name }}

conftest.py

@@ -14,6 +14,7 @@
 from devito.ir.iet import (FindNodes, FindSymbols, Iteration, ParallelBlock,
                            retrieve_iteration_tree)
 from devito.tools import as_tuple
+from devito.petsc.utils import get_petsc_dir, get_petsc_arch
 
 try:
     from mpi4py import MPI  # noqa
@@ -33,7 +34,7 @@ def skipif(items, whole_module=False):
     accepted = set()
     accepted.update({'device', 'device-C', 'device-openmp', 'device-openacc',
                      'device-aomp', 'cpu64-icc', 'cpu64-icx', 'cpu64-nvc', 'cpu64-arm',
-                     'cpu64-icpx', 'chkpnt'})
+                     'cpu64-icpx', 'chkpnt', 'petsc'})
     accepted.update({'nodevice'})
     unknown = sorted(set(items) - accepted)
     if unknown:
@@ -87,6 +88,19 @@ def skipif(items, whole_module=False):
         if i == 'chkpnt' and Revolver is NoopRevolver:
             skipit = "pyrevolve not installed"
             break
+        if i == 'petsc':
+            petsc_dir = get_petsc_dir()
+            petsc_arch = get_petsc_arch()
+            if petsc_dir is None or petsc_arch is None:
+                skipit = "PETSC_DIR or PETSC_ARCH are not set"
+                break
+            else:
+                petsc_installed = os.path.join(
+                    petsc_dir, petsc_arch, 'include', 'petscconf.h'
+                )
+                if not os.path.isfile(petsc_installed):
+                    skipit = "PETSc is not installed"
+                    break
 
     if skipit is False:
         return pytest.mark.skipif(False, reason='')

devito/ir/equations/algorithms.py

@@ -170,7 +170,8 @@ def concretize_subdims(exprs, **kwargs):
     """
     sregistry = kwargs.get('sregistry')
 
-    mapper = {}
+    # Update based on changes in #2509
+    mapper = kwargs.get('concretize_mapper', {})
     rebuilt = {}  # Rebuilt implicit dims etc which are shared between dimensions
 
     _concretize_subdims(exprs, mapper, rebuilt, sregistry)

devito/ir/equations/equation.py

@@ -9,10 +9,12 @@
                                Stencil, detect_io, detect_accesses)
 from devito.symbolics import IntDiv, limits_mapper, uxreplace
 from devito.tools import Pickable, Tag, frozendict
-from devito.types import Eq, Inc, ReduceMax, ReduceMin, relational_min
+from devito.types import (Eq, Inc, ReduceMax, ReduceMin,
+                          relational_min)
+from devito.types.equation import PetscEq
 
 __all__ = ['LoweredEq', 'ClusterizedEq', 'DummyEq', 'OpInc', 'OpMin', 'OpMax',
-           'identity_mapper']
+           'identity_mapper', 'OpPetsc']
 
 
 class IREq(sympy.Eq, Pickable):
@@ -102,7 +104,8 @@ def detect(cls, expr):
         reduction_mapper = {
             Inc: OpInc,
             ReduceMax: OpMax,
-            ReduceMin: OpMin
+            ReduceMin: OpMin,
+            PetscEq: OpPetsc
         }
         try:
             return reduction_mapper[type(expr)]
@@ -119,6 +122,7 @@ def detect(cls, expr):
 OpInc = Operation('+')
 OpMax = Operation('max')
 OpMin = Operation('min')
+OpPetsc = Operation('solve')
 
 
 identity_mapper = {

devito/ir/iet/algorithms.py

@@ -3,6 +3,8 @@
 from devito.ir.iet import (Expression, Increment, Iteration, List, Conditional, SyncSpot,
                            Section, HaloSpot, ExpressionBundle)
 from devito.tools import timed_pass
+from devito.petsc.types import MetaData
+from devito.petsc.iet.utils import petsc_iet_mapper
 
 __all__ = ['iet_build']
 
@@ -24,6 +26,8 @@ def iet_build(stree):
             for e in i.exprs:
                 if e.is_Increment:
                     exprs.append(Increment(e))
+                elif isinstance(e.rhs, MetaData):
+                    exprs.append(petsc_iet_mapper[e.operation](e, operation=e.operation))
                 else:
                     exprs.append(Expression(e, operation=e.operation))
             body = ExpressionBundle(i.ispace, i.ops, i.traffic, body=exprs)

devito/ir/iet/efunc.py

@@ -1,6 +1,6 @@
 from functools import cached_property
 
-from devito.ir.iet.nodes import Call, Callable
+from devito.ir.iet.nodes import Call, Callable, FixedArgsCallable
 from devito.ir.iet.utils import derive_parameters
 from devito.symbolics import uxreplace
 from devito.tools import as_tuple
@@ -131,7 +131,7 @@ class AsyncCall(Call):
     pass
 
 
-class ThreadCallable(Callable):
+class ThreadCallable(FixedArgsCallable):
 
     """
     A Callable executed asynchronously by a thread.

devito/ir/iet/nodes.py

@@ -1,6 +1,7 @@
 """The Iteration/Expression Tree (IET) hierarchy."""
 
 import abc
+import ctypes
 import inspect
 from functools import cached_property
 from collections import OrderedDict, namedtuple
@@ -19,7 +20,7 @@
                           ctypes_to_cstr)
 from devito.types.basic import (AbstractFunction, AbstractSymbol, Basic, Indexed,
                                 Symbol)
-from devito.types.object import AbstractObject, LocalObject
+from devito.types.object import AbstractObject, LocalObject, LocalCompositeObject
 
 __all__ = ['Node', 'MultiTraversable', 'Block', 'Expression', 'Callable',
            'Call', 'ExprStmt', 'Conditional', 'Iteration', 'List', 'Section',
@@ -28,7 +29,7 @@
            'Increment', 'Return', 'While', 'ListMajor', 'ParallelIteration',
            'ParallelBlock', 'Dereference', 'Lambda', 'SyncSpot', 'Pragma',
            'DummyExpr', 'BlankLine', 'ParallelTree', 'BusyWait', 'UsingNamespace',
-           'CallableBody', 'Transfer']
+           'CallableBody', 'Transfer', 'Callback', 'FixedArgsCallable']
 
 # First-class IET nodes
 
@@ -763,6 +764,15 @@ def defines(self):
         return self.all_parameters
 
 
+class FixedArgsCallable(Callable):
+
+    """
+    A Callable class that enforces a fixed function signature.
+    """
+
+    pass
+
+
 class CallableBody(MultiTraversable):
 
     """
@@ -1041,8 +1051,8 @@ class Dereference(ExprStmt, Node):
     The following cases are supported:
 
         * `pointer` is a PointerArray or TempFunction, and `pointee` is an Array.
-        * `pointer` is an ArrayObject representing a pointer to a C struct, and
-          `pointee` is a field in `pointer`.
+        * `pointer` is an ArrayObject or LocalCompositeObject representing a pointer
+           to a C struct, and `pointee` is a field in `pointer`.
     """
 
     is_Dereference = True
@@ -1061,13 +1071,15 @@ def functions(self):
 
     @property
     def expr_symbols(self):
-        ret = [self.pointer.indexed]
+        ret = []
         if self.pointer.is_PointerArray or self.pointer.is_TempFunction:
-            ret.append(self.pointee.indexed)
+            ret.extend([self.pointer.indexed, self.pointee.indexed])
             ret.extend(flatten(i.free_symbols for i in self.pointee.symbolic_shape[1:]))
             ret.extend(self.pointer.free_symbols)
         else:
-            ret.append(self.pointee._C_symbol)
+            assert (isinstance(self.pointer, LocalCompositeObject) or
+                    issubclass(self.pointer._C_ctype, ctypes._Pointer))
+            ret.extend([self.pointer._C_symbol, self.pointee._C_symbol])
         return tuple(filter_ordered(ret))
 
     @property
@@ -1133,6 +1145,45 @@ def defines(self):
         return tuple(self.parameters)
 
 
+class Callback(Call):
+    """
+    Base class for special callback types.
+
+    Parameters
+    ----------
+    name : str
+        The name of the callback.
+    retval : str
+        The return type of the callback.
+    param_types : str or list of str
+        The return type for each argument of the callback.
+
+    Notes
+    -----
+    - The reason Callback is an IET type rather than a SymPy type is
+    due to the fact that, when represented at the SymPy level, the IET
+    engine fails to bind the callback to a specific Call. Consequently,
+    errors occur during the creation of the call graph.
+    """
+    # TODO: Create a common base class for Call and Callback to avoid
+    # having arguments=None here
+    def __init__(self, name, retval=None, param_types=None, arguments=None):
+        super().__init__(name=name)
+        self.retval = retval
+        self.param_types = as_tuple(param_types)
+
+    @property
+    def callback_form(self):
+        """
+        A string representation of the callback form.
+
+        Notes
+        -----
+        To be overridden by subclasses.
+        """
+        return
+
+
 class Section(List):
 
     """

devito/ir/iet/visitors.py

@@ -24,7 +24,7 @@
                           c_restrict_void_p, sorted_priority)
 from devito.types.basic import AbstractFunction, Basic
 from devito.types import (ArrayObject, CompositeObject, Dimension, Pointer,
-                          IndexedData, DeviceMap)
+                          IndexedData, DeviceMap, LocalCompositeObject)
 
 
 __all__ = ['FindApplications', 'FindNodes', 'FindSections', 'FindSymbols',
@@ -186,11 +186,10 @@ def __init__(self, *args, compiler=None, **kwargs):
         '_mem_constant': 'static',
         '_mem_shared': '',
     }
-    _restrict_keyword = 'restrict'
 
     def _gen_struct_decl(self, obj, masked=()):
         """
-        Convert ctypes.Struct -> cgen.Structure.
+        Convert ctypes.Struct and LocalCompositeObject -> cgen.Structure.
         """
         ctype = obj._C_ctype
         try:
@@ -201,7 +200,16 @@ def _gen_struct_decl(self, obj, masked=()):
                 return None
         except TypeError:
             # E.g., `ctype` is of type `dtypes_lowering.CustomDtype`
-            return None
+            if isinstance(obj, LocalCompositeObject):
+                # TODO: Potentially re-evaluate: Setting ctype to obj allows
+                # _gen_struct_decl to generate a cgen.Structure from a
+                # LocalCompositeObject, where obj._C_ctype is a CustomDtype.
+                # LocalCompositeObject has a __fields__ property,
+                # which allows the subsequent code in this function to function
+                # correctly.
+                ctype = obj
+            else:
+                return None
 
         try:
             return obj._C_typedecl
@@ -250,10 +258,10 @@ def _gen_value(self, obj, mode=1, masked=()):
             strshape = ''
             if isinstance(obj, (AbstractFunction, IndexedData)) and mode >= 1:
                 if not obj._mem_stack:
-                    strtype = '%s%s' % (strtype, self._restrict_keyword)
+                    strtype = '%s%s' % (strtype, obj._restrict_keyword)
         strtype = ' '.join(qualifiers + [strtype])
 
-        if obj.is_LocalObject and obj._C_modifier is not None and mode == 2:
+        if obj.is_LocalType and obj._C_modifier is not None and mode == 2:
             strtype += obj._C_modifier
 
         strname = obj._C_name
@@ -624,6 +632,9 @@ def visit_Lambda(self, o):
                      (', '.join(captures), ', '.join(decls), ''.join(extra)))
         return LambdaCollection([top, c.Block(body)])
 
+    def visit_Callback(self, o, nested_call=False):
+        return CallbackArg(o)
+
     def visit_HaloSpot(self, o):
         body = flatten(self._visit(i) for i in o.children)
         return c.Collection(body)
@@ -675,8 +686,11 @@ def _operator_typedecls(self, o, mode='all'):
         for i in o._func_table.values():
             if not i.local:
                 continue
-            typedecls.extend([self._gen_struct_decl(j) for j in i.root.parameters
-                              if xfilter(j)])
+            typedecls.extend([
+                self._gen_struct_decl(j)
+                for j in FindSymbols().visit(i.root)
+                if xfilter(j)
+            ])
         typedecls = filter_sorted(typedecls, key=lambda i: i.tpname)
 
         return typedecls
@@ -1427,3 +1441,12 @@ def sorted_efuncs(efuncs):
         CommCallable: 1
     }
     return sorted_priority(efuncs, priority)
+
+
+class CallbackArg(c.Generable):
+
+    def __init__(self, callback):
+        self.callback = callback
+
+    def generate(self):
+        yield self.callback.callback_form