dsl: Add PETSc functionality

.github/workflows/examples-mpi.yml

@@ -17,9 +17,11 @@ on:
   push:
     branches:
       - main
+      - petsc
   pull_request:
     branches:
       - main
+      - petsc
 
 jobs:
   build:

.github/workflows/examples.yml

@@ -10,9 +10,11 @@ on:
   push:
     branches:
       - main
+      - petsc
   pull_request:
     branches:
       - main
+      - petsc
 
 jobs:
   tutorials:

.github/workflows/flake8.yml

@@ -10,9 +10,11 @@ on:
    push:
      branches:
        - main
+       - petsc
    pull_request:
      branches:
        - main
+       - petsc
 
 jobs:
   flake8:

.github/workflows/pytest-core-mpi.yml

@@ -10,9 +10,11 @@ on:
   push:
     branches:
       - main
+      - petsc
   pull_request:
     branches:
       - main
+      - petsc
 
 jobs:
   test-mpi-basic:

.github/workflows/pytest-core-nompi.yml

@@ -10,9 +10,11 @@ on:
   push:
     branches:
       - main
+      - petsc
   pull_request:
     branches:
       - main
+      - petsc
 
 jobs:
   pytest:

.github/workflows/pytest-petsc.yml

@@ -0,0 +1,97 @@
+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:
+      - main
+      - petsc
+  pull_request:
+    branches:
+      - main
+      - petsc
+
+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: Log in to DockerHub
+      uses: docker/login-action@v2
+      with:
+        username: ${{ secrets.DOCKER_USERNAME }}
+        password: ${{ secrets.DOCKER_PASSWORD }}
+
+    - name: Build docker image
+      run: |
+          docker build -f docker/Dockerfile.devito --build-arg base=zoeleibowitz/petsc_image:latest --tag zoeleibowitz/petsc_devito_image:latest .
+
+    - name: Set run prefix
+      run: |
+          echo "RUN_CMD=docker run --rm -t -e CODECOV_TOKEN=${{ secrets.CODECOV_TOKEN }} --name testrun zoeleibowitz/petsc_devito_image:latest"  >> $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 - serial
+      run: |
+        ${{ env.RUN_CMD }} mpiexec -n 1 pytest -m "not parallel" --cov --cov-config=.coveragerc --cov-report=xml ${{ env.TESTS }}
+
+    - name: Test with pytest - parallel
+      run: |
+        ${{ env.RUN_CMD }} mpiexec -n 1 pytest -m parallel --cov --cov-config=.coveragerc --cov-report=xml ${{ env.TESTS }}
+
+    - name: Test examples
+      run: |
+        ${{ env.RUN_CMD }} mpiexec -n 1 python3 examples/petsc/seismic/01_staggered_acoustic.py
+        ${{ env.RUN_CMD }} mpiexec -n 1 python3 examples/petsc/cfd/01_navierstokes.py
+        ${{ env.RUN_CMD }} mpiexec -n 1 python3 examples/petsc/Poisson/01_poisson.py
+        ${{ env.RUN_CMD }} mpiexec -n 1 python3 examples/petsc/Poisson/02_laplace.py
+        ${{ env.RUN_CMD }} mpiexec -n 1 python3 examples/petsc/Poisson/03_poisson.py
+        ${{ env.RUN_CMD }} mpiexec -n 1 python3 examples/petsc/Poisson/04_poisson.py
+        ${{ env.RUN_CMD }} mpiexec -n 1 python3 examples/petsc/random/01_helmholtz.py
+        ${{ env.RUN_CMD }} mpiexec -n 1 python3 examples/petsc/random/02_biharmonic.py
+
+    - name: Upload coverage to Codecov
+      if: "!contains(matrix.name, 'docker')"
+      uses: codecov/codecov-action@v4
+      with:
+        token: ${{ secrets.CODECOV_TOKEN }}
+        name: ${{ matrix.name }}

.github/workflows/tutorials.yml

@@ -10,9 +10,11 @@ on:
   push:
     branches:
       - main
+      - petsc
   pull_request:
     branches:
       - main
+      - petsc
 
 jobs:
   tutorials:

conftest.py

@@ -15,6 +15,7 @@
 from devito.ir.iet import (FindNodes, FindSymbols, Iteration, ParallelBlock,
                            retrieve_iteration_tree)
 from devito.tools import as_tuple
+from devito.petsc.config import PetscOSError, get_petsc_dir
 
 try:
     from mpi4py import MPI  # noqa
@@ -34,7 +35,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',
-                     'noadvisor', 'cpu64-arm', 'cpu64-icpx', 'chkpnt'})
+                     'noadvisor', 'cpu64-arm', 'cpu64-icpx', 'chkpnt', 'petsc'})
     accepted.update({'nodevice'})
     unknown = sorted(set(items) - accepted)
     if unknown:
@@ -94,6 +95,12 @@ def skipif(items, whole_module=False):
         if i == 'chkpnt' and Revolver is NoopRevolver:
             skipit = "pyrevolve not installed"
             break
+        if i == 'petsc':
+            try:
+                _ = get_petsc_dir()
+            except PetscOSError:
+                skipit = "PETSc is not installed"
+                break
 
     if skipit is False:
         return pytest.mark.skipif(False, reason='')

devito/core/__init__.py

@@ -7,6 +7,7 @@
     Cpu64CXXNoopCOperator, Cpu64CXXNoopOmpOperator,
     Cpu64AdvCXXOperator, Cpu64AdvCXXOmpOperator,
     Cpu64FsgCXXOperator, Cpu64FsgCXXOmpOperator,
+    Cpu64NoopPetscOperator, Cpu64AdvPetscOperator,
     Cpu64CustomCXXOmpOperator, Cpu64CustomCOperator
 )
 from devito.core.intel import (
@@ -45,12 +46,14 @@
 operator_registry.add(Cpu64NoopOmpOperator, Cpu64, 'noop', 'Copenmp')
 operator_registry.add(Cpu64CXXNoopCOperator, Cpu64, 'noop', 'CXX')
 operator_registry.add(Cpu64CXXNoopOmpOperator, Cpu64, 'noop', 'CXXopenmp')
+operator_registry.add(Cpu64NoopPetscOperator, Cpu64, 'noop', 'petsc')
 
 operator_registry.add(Cpu64AdvCOperator, Cpu64, 'advanced', 'C')
 operator_registry.add(Cpu64AdvOmpOperator, Cpu64, 'advanced', 'openmp')
 operator_registry.add(Cpu64AdvOmpOperator, Cpu64, 'advanced', 'Copenmp')
 operator_registry.add(Cpu64AdvCXXOperator, Cpu64, 'advanced', 'CXX')
 operator_registry.add(Cpu64AdvCXXOmpOperator, Cpu64, 'advanced', 'CXXopenmp')
+operator_registry.add(Cpu64AdvPetscOperator, Cpu64, 'advanced', 'petsc')
 
 operator_registry.add(Cpu64FsgCOperator, Cpu64, 'advanced-fsg', 'C')
 operator_registry.add(Cpu64FsgOmpOperator, Cpu64, 'advanced-fsg', 'openmp')

devito/core/cpu.py

@@ -11,13 +11,16 @@
 from devito.passes.iet import (CTarget, CXXTarget, COmpTarget, CXXOmpTarget,
                                avoid_denormals, linearize,
                                mpiize, hoist_prodders, relax_incr_dimensions,
-                               check_stability)
+                               check_stability, PetscTarget)
 from devito.tools import timed_pass
 
+from devito.petsc.iet.passes import lower_petsc_symbols
+
 __all__ = ['Cpu64NoopCOperator', 'Cpu64NoopOmpOperator', 'Cpu64AdvCOperator',
            'Cpu64AdvOmpOperator', 'Cpu64FsgCOperator', 'Cpu64FsgOmpOperator',
            'Cpu64CustomOperator', 'Cpu64CustomCXXOperator', 'Cpu64AdvCXXOperator',
-           'Cpu64AdvCXXOmpOperator', 'Cpu64FsgCXXOperator', 'Cpu64FsgCXXOmpOperator']
+           'Cpu64AdvCXXOmpOperator', 'Cpu64FsgCXXOperator', 'Cpu64FsgCXXOmpOperator',
+           'Cpu64NoopPetscOperator']
 
 
 class Cpu64OperatorMixin:
@@ -142,6 +145,9 @@ def _specialize_iet(cls, graph, **kwargs):
         # Symbol definitions
         cls._Target.DataManager(**kwargs).process(graph)
 
+        # Lower PETSc symbols
+        lower_petsc_symbols(graph, **kwargs)
+
         return graph
 
 
@@ -221,6 +227,9 @@ def _specialize_iet(cls, graph, **kwargs):
         # Symbol definitions
         cls._Target.DataManager(**kwargs).process(graph)
 
+        # Lower PETSc symbols
+        lower_petsc_symbols(graph, **kwargs)
+
         # Linearize n-dimensional Indexeds
         linearize(graph, **kwargs)
 
@@ -360,6 +369,15 @@ class Cpu64CXXNoopOmpOperator(Cpu64NoopOperator):
     LINEARIZE = True
 
 
+class Cpu64NoopPetscOperator(Cpu64NoopOperator):
+    _Target = PetscTarget
+
+    @classmethod
+    def _rcompile_wrapper(cls, **kwargs0):
+        kwargs0['language'] = 'petsc'
+        return super()._rcompile_wrapper(**kwargs0)
+
+
 class Cpu64AdvCOperator(Cpu64AdvOperator):
     _Target = CTarget
 
@@ -369,6 +387,15 @@ class Cpu64AdvCXXOperator(Cpu64AdvOperator):
     LINEARIZE = True
 
 
+class Cpu64AdvPetscOperator(Cpu64AdvOperator):
+    _Target = PetscTarget
+
+    @classmethod
+    def _rcompile_wrapper(cls, **kwargs0):
+        kwargs0['language'] = 'petsc'
+        return super()._rcompile_wrapper(**kwargs0)
+
+
 class Cpu64AdvOmpOperator(Cpu64AdvOperator):
     _Target = COmpTarget
 

devito/ir/equations/algorithms.py

@@ -191,7 +191,8 @@ def concretize_subdims(exprs, **kwargs):
     """
     sregistry = kwargs.get('sregistry')
 
-    mapper = {}
+    # To be updated 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.nodes 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.