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
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):
......
......@@ -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
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]
......
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)
......@@ -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))
......
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