Phasefield: prepared extraction of staggered quantities computation

- support for different #ghost layers at each coordinate limit

cpu/kernelcreation.py

@@ -6,7 +6,8 @@ from pystencils.field import Field
 import pystencils.ast as ast
 
 
-def createKernel(listOfEquations, functionName="kernel", typeForSymbol=None, splitGroups=(), iterationSlice=None):
+def createKernel(listOfEquations, functionName="kernel", typeForSymbol=None, splitGroups=(),
+                 iterationSlice=None, ghostLayers=None):
     """
     Creates an abstract syntax tree for a kernel function, by taking a list of update rules.
 
@@ -21,6 +22,10 @@ def createKernel(listOfEquations, functionName="kernel", typeForSymbol=None, spl
     :param splitGroups: Specification on how to split up inner loop into multiple loops. For details see
            transformation :func:`pystencils.transformation.splitInnerLoop`
     :param iterationSlice: if not None, iteration is done only over this slice of the field
+    :param ghostLayers: a sequence of pairs for each coordinate with lower and upper nr of ghost layers
+                        if None, the number of ghost layers is determined automatically and assumed to be equal for a
+                        all dimensions
+
     :return: :class:`pystencils.ast.KernelFunction` node
     """
     if not typeForSymbol:
@@ -43,7 +48,7 @@ def createKernel(listOfEquations, functionName="kernel", typeForSymbol=None, spl
         field.setReadOnly()
 
     body = ast.Block(assignments)
-    code = makeLoopOverDomain(body, functionName, iterationSlice=iterationSlice)
+    code = makeLoopOverDomain(body, functionName, iterationSlice=iterationSlice, ghostLayers=ghostLayers)
 
     if splitGroups:
         typedSplitGroups = [[typeSymbol(s) for s in splitGroup] for splitGroup in splitGroups]

field.py

@@ -156,6 +156,11 @@ class Field:
     def __repr__(self):
         return self._fieldName
 
+    def neighbor(self, coordId, offset):
+        offsetList = [0] * self.spatialDimensions
+        offsetList[coordId] = offset
+        return Field.Access(self, tuple(offsetList))
+
     def __getitem__(self, offset):
         if type(offset) is np.ndarray:
             offset = tuple(offset)

transformations.py

@@ -9,12 +9,15 @@ from pystencils.slicing import normalizeSlice
 import pystencils.ast as ast
 
 
-def makeLoopOverDomain(body, functionName, iterationSlice=None):
+def makeLoopOverDomain(body, functionName, iterationSlice=None, ghostLayers=None):
     """
     Uses :class:`pystencils.field.Field.Access` to create (multiple) loops around given AST.
     :param body: list of nodes
     :param functionName: name of generated C function
     :param iterationSlice: if not None, iteration is done only over this slice of the field
+    :param ghostLayers: a sequence of pairs for each coordinate with lower and upper nr of ghost layers
+                        if None, the number of ghost layers is determined automatically and assumed to be equal for a
+                        all dimensions
     :return: :class:`LoopOverCoordinate` instance with nested loops, ordered according to field layouts
     """
     # find correct ordering by inspecting participating FieldAccesses
@@ -23,9 +26,6 @@ def makeLoopOverDomain(body, functionName, iterationSlice=None):
     fields = set(fieldList)
     loopOrder = getOptimalLoopOrdering(fields)
 
-    # find number of required ghost layers
-    requiredGhostLayers = max([fa.requiredGhostLayers for fa in fieldAccesses])
-
     shapes = set([f.spatialShape for f in fields])
 
     if len(shapes) > 1:
@@ -39,12 +39,16 @@ def makeLoopOverDomain(body, functionName, iterationSlice=None):
     if iterationSlice is not None:
         iterationSlice = normalizeSlice(iterationSlice, shape)
 
+    if ghostLayers is None:
+        requiredGhostLayers = max([fa.requiredGhostLayers for fa in fieldAccesses])
+        ghostLayers = [(requiredGhostLayers, requiredGhostLayers)] * len(loopOrder)
+
     currentBody = body
     lastLoop = None
     for i, loopCoordinate in enumerate(loopOrder):
         if iterationSlice is None:
-            begin = requiredGhostLayers
-            end = shape[loopCoordinate] - requiredGhostLayers
+            begin = ghostLayers[loopCoordinate][0]
+            end = shape[loopCoordinate] - ghostLayers[loopCoordinate][1]
             newLoop = ast.LoopOverCoordinate(currentBody, loopCoordinate, begin, end, 1)
             lastLoop = newLoop
             currentBody = ast.Block([lastLoop])