Sliced iteration

- LoopOverDomain changed to be able to loop over rectangular sub-region of field
- support for slicing with makeSlice

__init__.py

 from pystencils.field import Field, extractCommonSubexpressions
-
+from pystencils.typedsymbol import TypedSymbol
+from pystencils.slicing import makeSlice

ast.py

@@ -118,6 +118,10 @@ class Block(Node):
         node.parent = self
         self._nodes.insert(0, node)
 
+    def insertBefore(self, newNode, insertBefore):
+        idx = self._nodes.index(insertBefore)
+        self._nodes.insert(idx, newNode)
+
     def append(self, node):
         node.parent = self
         self._nodes.append(node)
@@ -162,30 +166,26 @@ class PragmaBlock(Block):
 class LoopOverCoordinate(Node):
     LOOP_COUNTER_NAME_PREFIX = "ctr"
 
-    def __init__(self, body, coordinateToLoopOver, shape, increment=1, ghostLayers=1,
-                 isInnermostLoop=False, isOutermostLoop=False):
+    def __init__(self, body, coordinateToLoopOver, start, stop, step=1):
         self._body = body
         self._coordinateToLoopOver = coordinateToLoopOver
-        self._shape = shape
-        self._increment = increment
-        self._ghostLayers = ghostLayers
+        self._begin = start
+        self._end = stop
+        self._increment = step
         self._body.parent = self
         self.prefixLines = []
-        self._isInnermostLoop = isInnermostLoop
-        self._isOutermostLoop = isOutermostLoop
 
     def newLoopWithDifferentBody(self, newBody):
-        result = LoopOverCoordinate(newBody, self._coordinateToLoopOver, self._shape, self._increment,
-                                    self._ghostLayers, self._isInnermostLoop, self._isOutermostLoop)
+        result = LoopOverCoordinate(newBody, self._coordinateToLoopOver, self._begin, self._end, self._increment)
         result.prefixLines = self.prefixLines
         return result
 
     @property
     def args(self):
         result = [self._body]
-        limit = self._shape[self._coordinateToLoopOver]
-        if isinstance(limit, Node) or isinstance(limit, sp.Basic):
-            result.append(limit)
+        for e in [self._begin, self._end, self._increment]:
+            if hasattr(e, "args"):
+                result.append(e)
         return result
 
     @property
@@ -193,8 +193,16 @@ class LoopOverCoordinate(Node):
         return self._body
 
     @property
-    def iterationEnd(self):
-        return self._shape[self.coordinateToLoopOver] - self.ghostLayers
+    def start(self):
+        return self._begin
+
+    @property
+    def stop(self):
+        return self._end
+
+    @property
+    def step(self):
+        return self._increment
 
     @property
     def coordinateToLoopOver(self):
@@ -206,42 +214,44 @@ class LoopOverCoordinate(Node):
         result.add(self.loopCounterSymbol)
         return result
 
+    @staticmethod
+    def getLoopCounterName(coordinateToLoopOver):
+        return "%s_%s" % (LoopOverCoordinate.LOOP_COUNTER_NAME_PREFIX, coordinateToLoopOver)
+
     @property
     def loopCounterName(self):
-        return "%s_%s" % (LoopOverCoordinate.LOOP_COUNTER_NAME_PREFIX, self._coordinateToLoopOver)
+        return LoopOverCoordinate.getLoopCounterName(self.coordinateToLoopOver)
+
+    @staticmethod
+    def getLoopCounterSymbol(coordinateToLoopOver):
+        return TypedSymbol(LoopOverCoordinate.getLoopCounterName(coordinateToLoopOver), "int")
 
     @property
     def loopCounterSymbol(self):
-        return TypedSymbol(self.loopCounterName, "int")
+        return LoopOverCoordinate.getLoopCounterSymbol(self.coordinateToLoopOver)
 
     @property
     def symbolsRead(self):
-        result = self._body.symbolsRead
-        limit = self._shape[self._coordinateToLoopOver]
-        if isinstance(limit, sp.Basic):
-            result.update(limit.atoms(sp.Symbol))
+        loopBoundSymbols = set()
+        for possibleSymbol in [self._begin, self._end, self._increment]:
+            if isinstance(possibleSymbol, Node) or isinstance(possibleSymbol, sp.Basic):
+                loopBoundSymbols.update(possibleSymbol.atoms(sp.Symbol))
+        result = self._body.symbolsRead.union(loopBoundSymbols)
         return result
 
     @property
     def isOutermostLoop(self):
-        return self._isOutermostLoop
+        from pystencils.transformations import getNextParentOfType
+        return getNextParentOfType(self, LoopOverCoordinate) is None
 
     @property
     def isInnermostLoop(self):
-        return self._isInnermostLoop
+        return len(self.atoms(LoopOverCoordinate)) == 0
 
     @property
     def coordinateToLoopOver(self):
         return self._coordinateToLoopOver
 
-    @property
-    def iterationRegionWithGhostLayer(self):
-        return self._shape[self._coordinateToLoopOver]
-
-    @property
-    def ghostLayers(self):
-        return self._ghostLayers
-
 
 class SympyAssignment(Node):
 

backends/cbackend.py

@@ -105,9 +105,9 @@ class CBackend:
                     yield e
 
         counterVar = node.loopCounterName
-        start = "int %s = %d" % (counterVar, node.ghostLayers)
-        condition = "%s < %s" % (counterVar, self.sympyPrinter.doprint(node.iterationEnd))
-        update = "++%s" % (counterVar,)
+        start = "int %s = %s" % (counterVar, self.sympyPrinter.doprint(node.start))
+        condition = "%s < %s" % (counterVar, self.sympyPrinter.doprint(node.stop))
+        update = "%s += %s" % (counterVar, self.sympyPrinter.doprint(node.step),)
         loopStr = "for (%s; %s; %s)" % (start, condition, update)
         return LoopWithOptionalPrefix(loopStr, self._print(node.body), prefixLines=node.prefixLines)
 

cpu/kernelcreation.py

@@ -6,7 +6,7 @@ from pystencils.field import Field
 import pystencils.ast as ast
 
 
-def createKernel(listOfEquations, functionName="kernel", typeForSymbol=None, splitGroups=()):
+def createKernel(listOfEquations, functionName="kernel", typeForSymbol=None, splitGroups=(), iterationSlice=None):
     """
     Creates an abstract syntax tree for a kernel function, by taking a list of update rules.
 
@@ -20,6 +20,7 @@ def createKernel(listOfEquations, functionName="kernel", typeForSymbol=None, spl
            right hand side is a sympy Boolean which are assumed to be 'bool' .
     :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
     :return: :class:`pystencils.ast.KernelFunction` node
     """
     if not typeForSymbol:
@@ -42,7 +43,7 @@ def createKernel(listOfEquations, functionName="kernel", typeForSymbol=None, spl
         field.setReadOnly()
 
     body = ast.Block(assignments)
-    code = makeLoopOverDomain(body, functionName)
+    code = makeLoopOverDomain(body, functionName, iterationSlice=iterationSlice)
 
     if splitGroups:
         typedSplitGroups = [[typeSymbol(s) for s in splitGroup] for splitGroup in splitGroups]

slicing.py

+import sympy as sp
+
+class SliceMaker(object):
+    def __getitem__(self, item):
+        return item
+makeSlice = SliceMaker()
+
+
+def normalizeSlice(slices, sizes):
+    """Converts slices with floating point and/or negative entries to integer slices"""
+
+    if len(slices) != len(sizes):
+        raise ValueError("Slice dimension does not match sizes")
+
+    result = []
+
+    for s, size in zip(slices, sizes):
+        if type(s) is int:
+            result.append(s)
+            continue
+        if type(s) is float:
+            result.append(int(s * size))
+            continue
+
+        assert (type(s) is slice)
+
+        if s.start is None:
+            newStart = 0
+        elif type(s.start) is float:
+            newStart = int(s.start * size)
+        else:
+            newStart = s.start
+
+        if s.stop is None:
+            newStop = size
+        elif type(s.stop) is float:
+            newStop = int(s.stop * size)
+        elif not isinstance(s.stop, sp.Basic) and s.stop < 0:
+            newStop = size + s.stop
+        else:
+            newStop = s.stop
+
+        result.append(slice(newStart, newStop, s.step if s.step is not None else 1))
+
+    return tuple(result)

transformations.py

@@ -2,16 +2,19 @@ from collections import defaultdict
 import sympy as sp
 from sympy.logic.boolalg import Boolean
 from sympy.tensor import IndexedBase
+
 from pystencils.field import Field, offsetComponentToDirectionString
 from pystencils.typedsymbol import TypedSymbol
+from pystencils.slicing import normalizeSlice
 import pystencils.ast as ast
 
 
-def makeLoopOverDomain(body, functionName):
+def makeLoopOverDomain(body, functionName, iterationSlice=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
     :return: :class:`LoopOverCoordinate` instance with nested loops, ordered according to field layouts
     """
     # find correct ordering by inspecting participating FieldAccesses
@@ -33,13 +36,29 @@ def makeLoopOverDomain(body, functionName):
         assert nrOfFixedSizedFields <= 1, "Differently sized field accesses in loop body: " + str(shapes)
     shape = list(shapes)[0]
 
+    if iterationSlice is not None:
+        iterationSlice = normalizeSlice(iterationSlice, shape)
+
     currentBody = body
     lastLoop = None
     for i, loopCoordinate in enumerate(loopOrder):
-        newLoop = ast.LoopOverCoordinate(currentBody, loopCoordinate, shape, 1, requiredGhostLayers,
-                                         isInnermostLoop=(i == 0), isOutermostLoop=(i == len(loopOrder) - 1))
-        lastLoop = newLoop
-        currentBody = ast.Block([lastLoop])
+        if iterationSlice is None:
+            begin = requiredGhostLayers
+            end = shape[loopCoordinate] - requiredGhostLayers
+            newLoop = ast.LoopOverCoordinate(currentBody, loopCoordinate, begin, end, 1)
+            lastLoop = newLoop
+            currentBody = ast.Block([lastLoop])
+        else:
+            sliceComponent = iterationSlice[loopCoordinate]
+            if type(sliceComponent) is slice:
+                sc = sliceComponent
+                newLoop = ast.LoopOverCoordinate(currentBody, loopCoordinate, sc.start, sc.stop, sc.step)
+                lastLoop = newLoop
+                currentBody = ast.Block([lastLoop])
+            else:
+                assignment = ast.SympyAssignment(ast.LoopOverCoordinate.getLoopCounterSymbol(loopCoordinate),
+                                                 sp.sympify(sliceComponent))
+                currentBody.insertFront(assignment)
     return ast.KernelFunction(currentBody, functionName)
 
 
@@ -236,21 +255,28 @@ def moveConstantsBeforeLoop(astNode):
     :return:
     """
     def findBlockToMoveTo(node):
-        """Traverses parents of node as long as the symbols are independent and returns a (parent) block
+        """
+        Traverses parents of node as long as the symbols are independent and returns a (parent) block
         the assignment can be safely moved to
-        :param node: SympyAssignment inside a Block"""
+        :param node: SympyAssignment inside a Block
+        :return blockToInsertTo, childOfBlockToInsertBefore
+        """
         assert isinstance(node, ast.SympyAssignment)
         assert isinstance(node.parent, ast.Block)
 
         lastBlock = node.parent
+        lastBlockChild = node
         element = node.parent
+        prevElement = node
         while element:
             if isinstance(element, ast.Block):
                 lastBlock = element
+                lastBlockChild = prevElement
             if node.symbolsRead.intersection(element.symbolsDefined):
                 break
+            prevElement = element
             element = element.parent
-        return lastBlock
+        return lastBlock, lastBlockChild
 
     def checkIfAssignmentAlreadyInBlock(assignment, targetBlock):
         for arg in targetBlock.args:
@@ -266,13 +292,13 @@ def moveConstantsBeforeLoop(astNode):
             if not isinstance(child, ast.SympyAssignment):
                 block.append(child)
             else:
-                target = findBlockToMoveTo(child)
+                target, childToInsertBefore = findBlockToMoveTo(child)
                 if target == block:     # movement not possible
                     target.append(child)
                 else:
                     existingAssignment = checkIfAssignmentAlreadyInBlock(child, target)
                     if not existingAssignment:
-                        target.insertFront(child)
+                        target.insertBefore(child, childToInsertBefore)
                     else:
                         assert existingAssignment.rhs == child.rhs, "Symbol with same name exists already"
 
@@ -335,7 +361,7 @@ def splitInnerLoop(astNode, symbolGroups):
     innerLoop.parent.replace(innerLoop, newLoops)
 
     for tmpArray in symbolsWithTemporaryArray:
-        outerLoop.parent.insertFront(ast.TemporaryMemoryAllocation(tmpArray, innerLoop.iterationRegionWithGhostLayer))
+        outerLoop.parent.insertFront(ast.TemporaryMemoryAllocation(tmpArray, innerLoop.stop))
         outerLoop.parent.append(ast.TemporaryMemoryFree(tmpArray))