Utilise PetscSection to constrain essential boundary conditions

.github/workflows/docker-petsc.yml

@@ -6,7 +6,9 @@ permissions:
 on:
   push:
     branches:
-      - petsc         # Push events on petsc branch
+      - petsc
+      # will be dropped after merge into petsc branch
+      - petscsection
 
 jobs:
   build-and-push:
@@ -16,6 +18,9 @@ jobs:
       # Use buildkit
       DOCKER_BUILDKIT: "1"
 
+      PETSC_REPO: https://gitlab.com/ZoeLeibowitz/petsc.git
+      PETSC_BRANCH: zoe/feature-da-section-sf
+
     steps:
       - name: Checkout devito
         uses: actions/checkout@v5
@@ -38,10 +43,13 @@ jobs:
           context: .
           file: docker/Dockerfile.petsc
           push: true
+          platforms: linux/amd64
           tags: |
             devitocodes/devito-petsc:latest
-          build-args: base=devitocodes/devito:gcc-dev-amd64
-          platforms: linux/amd64
+          build-args: |
+            base=devitocodes/devito:gcc-dev-amd64
+            PETSC_REPO=${{ env.PETSC_REPO }}
+            PETSC_BRANCH=${{ env.PETSC_BRANCH }}
 
       - name: Remove dangling layers
         run: docker system prune -f

.github/workflows/pytest-petsc.yml

@@ -5,8 +5,6 @@ concurrency:
   cancel-in-progress: true
 
 on:
-  # Trigger the workflow on push or pull request,
-  # but only for the master branch
   push:
     branches:
       - main
@@ -50,7 +48,11 @@ jobs:
 
     - name: Build docker image
       run: |
-          docker build -f docker/Dockerfile.petsc --tag devito_petsc_image:test .
+          docker build \
+            -f docker/Dockerfile.petsc \
+            --build-arg PETSC_REPO=https://gitlab.com/ZoeLeibowitz/petsc.git \
+            --build-arg PETSC_BRANCH=zoe/feature-da-section-sf \
+            --tag devito_petsc_image:test .
 
     - name: Set run prefix
       run: |

devito/data/decomposition.py

@@ -450,6 +450,39 @@ def index_loc_to_glb(self, *args):
         else:
             raise TypeError("Expected 1 arguments, found %d" % len(args))
 
+    def index_glb_to_loc_unsafe(self, glb_idx, rel=True):
+        """
+        Convert a global index to a local index even if not owned.
+        WARNING: Must not be used to index data as there are no guard
+        rails against returning out of bound indices.
+        """
+        if not self.loc_empty:
+            loc_abs_min = self.loc_abs_min - self.glb_min
+            loc_abs_max = self.loc_abs_max - self.glb_min
+            glb_max = self.glb_max - self.glb_min
+        else:
+            loc_abs_min = self.loc_abs_min
+            loc_abs_max = self.loc_abs_max
+            glb_max = self.glb_max
+
+        glb_min = 0
+
+        base = loc_abs_min if rel else 0
+
+        # index_glb_to_loc(index)
+        # -> Base case, empty local subdomain
+        if self.loc_empty:
+            return None
+        # -> Handle negative index
+        if glb_idx < 0:
+            glb_idx = glb_max + glb_idx + 1
+        # -> Do the actual conversion
+        if loc_abs_min <= glb_idx <= loc_abs_max or glb_min <= glb_idx <= glb_max:
+            return glb_idx - base
+        else:
+            # This should raise an exception when used to access a numpy.array
+            return glb_idx
+
     def reshape(self, *args):
         """
         Create a new Decomposition with extended or reduced boundary subdomains.

devito/ir/cgen/printer.py

@@ -18,6 +18,7 @@
 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.misc import PostIncrementIndex
 from devito.tools import ctypes_to_cstr, dtype_to_ctype, ctypes_vector_mapper
 
 __all__ = ['BasePrinter', 'ccode']
@@ -148,8 +149,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):
         """

devito/ir/equations/equation.py

@@ -107,10 +107,10 @@ def detect(cls, expr):
             ReduceMin: OpMin,
             PetscEq: OpPetsc
         }
-        try:
-            return reduction_mapper[type(expr)]
-        except KeyError:
-            pass
+
+        for expr_type, op in reduction_mapper.items():
+            if isinstance(expr, expr_type):
+                return op
 
         # NOTE: in the future we might want to track down other kinds
         # of operations here (e.g., memcpy). However, we don't care for

devito/ir/iet/nodes.py

@@ -1148,8 +1148,8 @@ def expr_symbols(self):
                 ret.extend([self.pointer._C_symbol, self.pointee._C_symbol])
             else:
                 ret.extend([self.pointer, self.pointee.indexed])
-            ret.extend(flatten(i.free_symbols
-                               for i in self.pointee.symbolic_shape[1:]))
+                ret.extend(flatten(i.free_symbols
+                           for i in self.pointee.symbolic_shape[1:]))
         else:
             assert False, f"Unexpected pointer type {type(self.pointer)}"
 

devito/operator/operator.py

@@ -42,6 +42,7 @@
 from devito.types.dimension import Thickness
 from devito.petsc.iet.passes import lower_petsc
 from devito.petsc.clusters import petsc_preprocess
+from devito.petsc.equations import lower_exprs_petsc
 
 __all__ = ['Operator']
 
@@ -368,6 +369,8 @@ def _lower_exprs(cls, expressions, **kwargs):
         # in particular uniqueness across expressions is ensured
         expressions = concretize_subdims(expressions, **kwargs)
 
+        expressions = lower_exprs_petsc(expressions, **kwargs)
+
         processed = [LoweredEq(i) for i in expressions]
 
         return processed

devito/petsc/config.py

@@ -44,7 +44,8 @@ def core_metadata():
     petsc_lib = tuple([arch / 'lib' for arch in petsc_dir])
 
     return {
-        'includes': ('petscsnes.h', 'petscdmda.h'),
+        # TODO: Only add petscsection header when needed
+        'includes': ('petscsnes.h', 'petscdmda.h', 'petscsection.h'),
         'include_dirs': petsc_include,
         'libs': ('petsc'),
         'lib_dirs': petsc_lib,

devito/petsc/equations.py

@@ -0,0 +1,100 @@
+from devito.symbolics import retrieve_indexed, retrieve_dimensions, uxreplace
+from devito.types.dimension import SpaceDimension, CustomDimension
+from devito import Min, Max
+
+from devito.petsc.types.equation import ConstrainBC
+from devito.petsc.types.dimension import (
+    SubDimMax, SubDimMin,
+    SpaceDimMax, SpaceDimMin,
+)
+
+
+def lower_exprs_petsc(expressions, **kwargs):
+
+    # Process `ConstrainBC` equations
+    expressions = constrain_essential_bcs(expressions, **kwargs)
+
+    return expressions
+
+
+def constrain_essential_bcs(expressions, **kwargs):
+    """
+    """
+    constrain_expressions = [e for e in expressions if isinstance(e, ConstrainBC)]
+    if not constrain_expressions:
+        return expressions
+
+    sregistry = kwargs.get('sregistry')
+    new_exprs = []
+
+    # TODO: rethink
+    halo_size = {e.target.function._size_halo for e in constrain_expressions}
+    assert len(halo_size) == 1
+    halo_size = halo_size.pop()
+
+    all_dims = {d for e in constrain_expressions for d in extract_dims(e)}
+    subdims = [d for d in all_dims if d.is_Sub and not d.local]
+    space_dims = [d for d in all_dims if isinstance(d, SpaceDimension)]
+
+    mapper = {}
+
+    for d in subdims:
+        halo = halo_size[d]
+
+        subdim_max = SubDimMax(
+            sregistry.make_name(prefix=f"{d.name}_max"), subdim=d
+        )
+        subdim_min = SubDimMin(
+            sregistry.make_name(prefix=f"{d.name}_min"), subdim=d
+        )
+
+        mapper[d] = CustomDimension(
+            name=d.name,
+            symbolic_min=Max(subdim_min, d.parent.symbolic_min - halo.left),
+            symbolic_max=Min(subdim_max, d.parent.symbolic_max + halo.right),
+        )
+
+    for d in space_dims:
+        halo = halo_size[d]
+        space_dim_max = SpaceDimMax(
+            sregistry.make_name(prefix=f"{d.name}_max"), space_dim=d
+        )
+        space_dim_min = SpaceDimMin(
+            sregistry.make_name(prefix=f"{d.name}_min"), space_dim=d
+        )
+
+        mapper[d] = CustomDimension(
+            name=sregistry.make_name(prefix=f"{d.name}_expanded"),
+            symbolic_min=Max(space_dim_min, d.symbolic_min - halo.left),
+            symbolic_max=Min(space_dim_max, d.symbolic_max + halo.right),
+        )
+
+    # Apply mapper to expressions
+    for e in expressions:
+        if not isinstance(e, ConstrainBC):
+            new_exprs.append(e)
+            continue
+
+        dims = extract_dims(e)
+        if not dims:
+            new_exprs.append(e)
+            continue
+
+        new_e = uxreplace(e, mapper)
+
+        if e.implicit_dims:
+            new_e = new_e._rebuild(
+                implicit_dims=tuple(mapper.get(d, d) for d in e.implicit_dims)
+            )
+        new_exprs.append(new_e)
+    return new_exprs
+
+
+def extract_dims(expr):
+    indexeds = retrieve_indexed(expr)
+    dims = retrieve_dimensions(
+        [i for j in indexeds for i in j.indices],
+        mode="unique",
+    )
+    dims.update(expr.implicit_dims)
+    return dims