Commit 411af476 authored by Martin Bauer's avatar Martin Bauer
Browse files

Sliced iteration

- LoopOverDomain changed to be able to loop over rectangular sub-region of field
- support for slicing with makeSlice
parent 3e00ddfd
from pystencils.field import Field, extractCommonSubexpressions
from pystencils.typedsymbol import TypedSymbol
from pystencils.slicing import makeSlice
......@@ -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
......@@ -162,30 +166,26 @@ class PragmaBlock(Block):
class LoopOverCoordinate(Node):
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
def args(self):
result = [self._body]
limit = self._shape[self._coordinateToLoopOver]
if isinstance(limit, Node) or isinstance(limit, sp.Basic):
for e in [self._begin, self._end, self._increment]:
if hasattr(e, "args"):
return result
......@@ -193,8 +193,16 @@ class LoopOverCoordinate(Node):
return self._body
def iterationEnd(self):
return self._shape[self.coordinateToLoopOver] - self.ghostLayers
def start(self):
return self._begin
def stop(self):
return self._end
def step(self):
return self._increment
def coordinateToLoopOver(self):
......@@ -206,42 +214,44 @@ class LoopOverCoordinate(Node):
return result
def getLoopCounterName(coordinateToLoopOver):
return "%s_%s" % (LoopOverCoordinate.LOOP_COUNTER_NAME_PREFIX, coordinateToLoopOver)
def loopCounterName(self):
return "%s_%s" % (LoopOverCoordinate.LOOP_COUNTER_NAME_PREFIX, self._coordinateToLoopOver)
return LoopOverCoordinate.getLoopCounterName(self.coordinateToLoopOver)
def getLoopCounterSymbol(coordinateToLoopOver):
return TypedSymbol(LoopOverCoordinate.getLoopCounterName(coordinateToLoopOver), "int")
def loopCounterSymbol(self):
return TypedSymbol(self.loopCounterName, "int")
return LoopOverCoordinate.getLoopCounterSymbol(self.coordinateToLoopOver)
def symbolsRead(self):
result = self._body.symbolsRead
limit = self._shape[self._coordinateToLoopOver]
if isinstance(limit, sp.Basic):
loopBoundSymbols = set()
for possibleSymbol in [self._begin, self._end, self._increment]:
if isinstance(possibleSymbol, Node) or isinstance(possibleSymbol, sp.Basic):
result = self._body.symbolsRead.union(loopBoundSymbols)
return result
def isOutermostLoop(self):
return self._isOutermostLoop
from pystencils.transformations import getNextParentOfType
return getNextParentOfType(self, LoopOverCoordinate) is None
def isInnermostLoop(self):
return self._isInnermostLoop
return len(self.atoms(LoopOverCoordinate)) == 0
def coordinateToLoopOver(self):
return self._coordinateToLoopOver
def iterationRegionWithGhostLayer(self):
return self._shape[self._coordinateToLoopOver]
def ghostLayers(self):
return self._ghostLayers
class SympyAssignment(Node):
......@@ -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)
......@@ -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
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]
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:
if type(s) is float:
result.append(int(s * size))
assert (type(s) is slice)
if s.start is None:
newStart = 0
elif type(s.start) is float:
newStart = int(s.start * size)
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
newStop = s.stop
result.append(slice(newStart, newStop, s.step if s.step is not None else 1))
return tuple(result)
......@@ -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])
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])
assignment = ast.SympyAssignment(ast.LoopOverCoordinate.getLoopCounterSymbol(loopCoordinate),
return ast.KernelFunction(currentBody, functionName)
......@@ -236,21 +255,28 @@ def moveConstantsBeforeLoop(astNode):
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):
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):
target = findBlockToMoveTo(child)
target, childToInsertBefore = findBlockToMoveTo(child)
if target == block: # movement not possible
existingAssignment = checkIfAssignmentAlreadyInBlock(child, target)
if not existingAssignment:
target.insertBefore(child, childToInsertBefore)
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))
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment