diff --git a/__init__.py b/__init__.py
index 8e21541ca74e622cf6ae7420c5b339fcde2a4610..2f6f6bf2f64761f6127ce91e7c4c0c28e1e73ee2 100644
--- a/__init__.py
+++ b/__init__.py
@@ -1,2 +1,3 @@
from pystencils.field import Field, extractCommonSubexpressions
-
+from pystencils.typedsymbol import TypedSymbol
+from pystencils.slicing import makeSlice
diff --git a/ast.py b/ast.py
index 301bf81cff3b11280c566a33d1c3c2880ada1190..7600e946f0195a6d30722970d9b0030182c598c8 100644
--- a/ast.py
+++ b/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):
diff --git a/backends/cbackend.py b/backends/cbackend.py
index ef24a39fb3b389f08ba62e6c2685de5122b25720..bf5f92a06e24dc55025d49a3165e6a7380fa47b7 100644
--- a/backends/cbackend.py
+++ b/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)
diff --git a/cpu/kernelcreation.py b/cpu/kernelcreation.py
index 185ebaa3c759a50c9fb230d1cee3f50b236d5a7e..6407382bae6f7ef64af77a354f009eccc8f8c6bb 100644
--- a/cpu/kernelcreation.py
+++ b/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]
diff --git a/slicing.py b/slicing.py
new file mode 100644
index 0000000000000000000000000000000000000000..c9b45a6592531921b8d3522a68ac22528a73418a
--- /dev/null
+++ b/slicing.py
@@ -0,0 +1,45 @@
+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)
diff --git a/transformations.py b/transformations.py
index 6ea7014fb98c5e00dfe0f55621cc2280de89ffdc..1443908037d4b025d201bdd3893e487a11c50032 100644
--- a/transformations.py
+++ b/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))