Commit aea202a7 authored by Martin Bauer's avatar Martin Bauer
Browse files

Documentation & Restructuring

- added sphinx files for documentation generation
- collected kernel creation functions in new "cpu" and "cudagpu" modules
parent bd56d31f
from pystencils.field import Field, extractCommonSubexpressions
......@@ -5,21 +5,29 @@ from pystencils.typedsymbol import TypedSymbol
class Node:
"""Base class for all AST nodes"""
def __init__(self, parent=None):
self.parent = parent
def args(self):
"""Returns all arguments/children of this node"""
return []
@property
def symbolsDefined(self):
"""Set of symbols which are defined in this node or its children"""
return set()
@property
def symbolsRead(self):
"""Set of symbols which are accessed/read in this node or its children"""
return set()
def atoms(self, argType):
"""
Returns a set of all children which are an instance of the given argType
"""
result = set()
for arg in self.args:
if isinstance(arg, argType):
......
......@@ -4,6 +4,9 @@ from pystencils.ast import Node
def printCCode(astNode):
"""
Prints the abstract syntax tree as C function
"""
printer = CBackend(cuda=False)
return printer(astNode)
......
from pystencils.cpu.kernelcreation import createKernel, addOpenMP
from pystencils.cpu.cpujit import makePythonFunction
from pystencils.backends.cbackend import printCCode
......@@ -131,6 +131,18 @@ def makePythonFunctionIncompleteParams(kernelFunctionNode, argumentDict):
def makePythonFunction(kernelFunctionNode, argumentDict={}):
"""
Creates C code from the abstract syntax tree, compiles it and makes it accessible as Python function
The parameters of the kernel are:
- numpy arrays for each field used in the kernel. The keyword argument name is the name of the field
- all symbols which are not defined in the kernel itself are expected as parameters
:param kernelFunctionNode: the abstract syntax tree
:param argumentDict: parameters passed here are already fixed. Remaining parameters have to be passed to the
returned kernel functor.
:return: kernel functor
"""
# build up list of CType arguments
try:
args = buildCTypeArgumentList(kernelFunctionNode, argumentDict)
......@@ -140,10 +152,3 @@ def makePythonFunction(kernelFunctionNode, argumentDict={}):
func = compileAndLoad(kernelFunctionNode)[kernelFunctionNode.functionName]
func.restype = None
return lambda: func(*args)
import sympy as sp
from pystencils.transformations import resolveFieldAccesses, makeLoopOverDomain, typingFromSympyInspection, \
typeAllEquations, getOptimalLoopOrdering, parseBasePointerInfo, moveConstantsBeforeLoop, splitInnerLoop
from pystencils.typedsymbol import TypedSymbol
from pystencils.field import Field
import pystencils.ast as ast
def createKernel(listOfEquations, functionName="kernel", typeForSymbol=None, splitGroups=()):
"""
Creates an abstract syntax tree for a kernel function, by taking a list of update rules.
Loops are created according to the field accesses in the equations.
:param listOfEquations: list of sympy equations, containing accesses to :class:`pystencils.field.Field`.
Defining the update rules of the kernel
:param functionName: name of the generated function - only important if generated code is written out
:param typeForSymbol: a map from symbol name to a C type specifier. If not specified all symbols are assumed to
be of type 'double' except symbols which occur on the left hand side of equations where the
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`
:return: :class:`pystencils.ast.KernelFunction` node
"""
if not typeForSymbol:
typeForSymbol = typingFromSympyInspection(listOfEquations, "double")
def typeSymbol(term):
if isinstance(term, Field.Access) or isinstance(term, TypedSymbol):
return term
elif isinstance(term, sp.Symbol):
return TypedSymbol(term.name, typeForSymbol[term.name])
else:
raise ValueError("Term has to be field access or symbol")
fieldsRead, fieldsWritten, assignments = typeAllEquations(listOfEquations, typeForSymbol)
allFields = fieldsRead.union(fieldsWritten)
for field in allFields:
field.setReadOnly(False)
for field in fieldsRead - fieldsWritten:
field.setReadOnly()
body = ast.Block(assignments)
code = makeLoopOverDomain(body, functionName)
if splitGroups:
typedSplitGroups = [[typeSymbol(s) for s in splitGroup] for splitGroup in splitGroups]
splitInnerLoop(code, typedSplitGroups)
loopOrder = getOptimalLoopOrdering(allFields)
basePointerInfo = [['spatialInner0'], ['spatialInner1']]
basePointerInfos = {field.name: parseBasePointerInfo(basePointerInfo, loopOrder, field) for field in allFields}
resolveFieldAccesses(code, fieldToBasePointerInfo=basePointerInfos)
moveConstantsBeforeLoop(code)
return code
def addOpenMP(astNode, schedule="static"):
"""
Parallelizes the outer loop with OpenMP
:param astNode: abstract syntax tree created e.g. by :func:`createKernel`
:param schedule: OpenMP scheduling policy e.g. 'static' or 'dynamic'
"""
assert type(astNode) is ast.KernelFunction
body = astNode.body
wrapperBlock = ast.PragmaBlock('#pragma omp parallel', body.takeChildNodes())
body.append(wrapperBlock)
outerLoops = [l for l in body.atoms(ast.LoopOverCoordinate) if l.isOutermostLoop]
assert outerLoops, "No outer loop found"
assert len(outerLoops) <= 1, "More than one outer loop found. Which one should be parallelized?"
outerLoops[0].prefixLines.append("#pragma omp for schedule(%s)" % (schedule,))
......@@ -6,147 +6,31 @@ from sympy.tensor import IndexedBase
from pystencils.typedsymbol import TypedSymbol
def getLayoutFromNumpyArray(arr):
"""
Returns a list indicating the memory layout (linearization order) of the numpy array.
Example:
>>> getLayoutFromNumpyArray(np.zeros([3,3,3]))
[0, 1, 2]
In this example the loop over the zeroth coordinate should be the outermost loop,
followed by the first and second. Elements arr[x,y,0] and arr[x,y,1] are adjacent in memory.
Normally constructed numpy arrays have this order, however by stride tricks or other frameworks, arrays
with different memory layout can be created.
"""
coordinates = list(range(len(arr.shape)))
return [x for (y, x) in sorted(zip(arr.strides, coordinates), key=lambda pair: pair[0], reverse=True)]
def numpyDataTypeToC(dtype):
"""Mapping numpy data types to C data types"""
if dtype == np.float64:
return "double"
elif dtype == np.float32:
return "float"
elif dtype == np.int32:
return "int"
raise NotImplementedError()
def offsetComponentToDirectionString(coordinateId, value):
"""
Translates numerical offset to string notation.
x offsets are labeled with east 'E' and 'W',
y offsets with north 'N' and 'S' and
z offsets with top 'T' and bottom 'B'
If the absolute value of the offset is bigger than 1, this number is prefixed.
:param coordinateId: integer 0, 1 or 2 standing for x,y and z
:param value: integer offset
Example:
>>> offsetComponentToDirectionString(0, 1)
'E'
>>> offsetComponentToDirectionString(1, 2)
'2N'
"""
nameComponents = (('W', 'E'), # west, east
('S', 'N'), # south, north
('B', 'T'), # bottom, top
)
if value == 0:
result = ""
elif value < 0:
result = nameComponents[coordinateId][0]
else:
result = nameComponents[coordinateId][1]
if abs(value) > 1:
result = "%d%s" % (abs(value), result)
return result
def offsetToDirectionString(offsetTuple):
"""
Translates numerical offset to string notation.
For details see :func:`offsetComponentToDirectionString`
:param offsetTuple: 3-tuple with x,y,z offset
Example:
>>> offsetToDirectionString([1, -1, 0])
'SE'
>>> offsetToDirectionString(([-3, 0, -2]))
'2B3W'
"""
names = ["", "", ""]
for i in range(len(offsetTuple)):
names[i] = offsetComponentToDirectionString(i, offsetTuple[i])
name = "".join(reversed(names))
if name == "":
name = "C"
return name
def directionStringToOffset(directionStr, dim=3):
"""
Reverse mapping of :func:`offsetToDirectionString`
:param directionStr: string representation of offset
:param dim: dimension of offset, i.e the length of the returned list
>>> directionStringToOffset('NW', dim=3)
array([-1, 1, 0])
>>> directionStringToOffset('NW', dim=2)
array([-1, 1])
>>> directionStringToOffset(offsetToDirectionString([3,-2,1]))
array([ 3, -2, 1])
"""
offsetMap = {
'C': np.array([0, 0, 0]),
'W': np.array([-1, 0, 0]),
'E': np.array([1, 0, 0]),
'S': np.array([0, -1, 0]),
'N': np.array([0, 1, 0]),
'B': np.array([0, 0, -1]),
'T': np.array([0, 0, 1]),
}
offset = np.array([0, 0, 0])
while len(directionStr) > 0:
factor = 1
firstNonDigit = 0
while directionStr[firstNonDigit].isdigit():
firstNonDigit += 1
if firstNonDigit > 0:
factor = int(directionStr[:firstNonDigit])
directionStr = directionStr[firstNonDigit:]
curOffset = offsetMap[directionStr[0]]
offset += factor * curOffset
directionStr = directionStr[1:]
return offset[:dim]
class Field:
"""
With fields one can formulate stencil-like update rules on structured grids.
This Field class knows about the dimension, memory layout (strides) and optionally about the size of an array.
To create a field use one of the static create* members. There are two options:
Creating Fields:
To create a field use one of the static create* members. There are two options:
1. create a kernel with fixed loop sizes i.e. the shape of the array is already known. This is usually the
case if just-in-time compilation directly from Python is done. (see Field.createFromNumpyArray)
case if just-in-time compilation directly from Python is done. (see :func:`Field.createFromNumpyArray`)
2. create a more general kernel that works for variable array sizes. This can be used to create kernels
beforehand for a library. (see Field.createGeneric)
beforehand for a library. (see :func:`Field.createGeneric`)
Dimensions:
A field has spatial and index dimensions, where the spatial dimensions come first.
The interpretation is that the field has multiple cells in (usually) two or three dimensional space which are
looped over. Additionally N values are stored per cell. In this case spatialDimensions is two or three,
and indexDimensions equals N. If you want to store a matrix on each point in a two dimensional grid, there
are four dimensions, two spatial and two index dimensions. len(arr.shape) == spatialDims + indexDims
are four dimensions, two spatial and two index dimensions: ``len(arr.shape) == spatialDims + indexDims``
Indexing:
When accessing (indexing) a field the result is a FieldAccess which is derived from sympy Symbol.
First specify the spatial offsets in [], then in case indexDimension>0 the indices in ()
e.g. f[-1,0,0](7)
e.g. ``f[-1,0,0](7)``
Example without index dimensions:
>>> a = np.zeros([10, 10])
......@@ -162,6 +46,7 @@ class Field:
Eq(dst_C^0, src_C^0)
Eq(dst_C^1, src_S^1)
Eq(dst_C^2, src_N^2)
"""
@staticmethod
def createFromNumpyArray(fieldName, npArray, indexDimensions=0):
......@@ -194,8 +79,8 @@ class Field:
:param spatialDimensions: see documentation of Field
:param indexDimensions: see documentation of Field
:param layout: tuple specifying the loop ordering of the spatial dimensions e.g. (2, 1, 0 ) means that
the outer loop loops over dimension 2, the second outer over dimension 1, and the inner loop
over dimension 0
the outer loop loops over dimension 2, the second outer over dimension 1, and the inner loop
over dimension 0
"""
if not layout:
layout = tuple(reversed(range(spatialDimensions)))
......@@ -384,3 +269,137 @@ class Field:
superClassContents = list(super(Field.Access, self)._hashable_content())
t = tuple([*superClassContents, hash(self._field), self._index] + self._offsets)
return t
def extractCommonSubexpressions(equations):
"""
Uses sympy to find common subexpressions in equations and returns
them in a topologically sorted order, ready for evaluation.
Usually called before list of equations is passed to :func:`createKernel`
"""
replacements, newEq = sp.cse(equations)
replacementEqs = [sp.Eq(*r) for r in replacements]
equations = replacementEqs + newEq
topologicallySortedPairs = sp.cse_main.reps_toposort([[e.lhs, e.rhs] for e in equations])
equations = [sp.Eq(*a) for a in topologicallySortedPairs]
return equations
def getLayoutFromNumpyArray(arr):
"""
Returns a list indicating the memory layout (linearization order) of the numpy array.
Example:
>>> getLayoutFromNumpyArray(np.zeros([3,3,3]))
[0, 1, 2]
In this example the loop over the zeroth coordinate should be the outermost loop,
followed by the first and second. Elements arr[x,y,0] and arr[x,y,1] are adjacent in memory.
Normally constructed numpy arrays have this order, however by stride tricks or other frameworks, arrays
with different memory layout can be created.
"""
coordinates = list(range(len(arr.shape)))
return [x for (y, x) in sorted(zip(arr.strides, coordinates), key=lambda pair: pair[0], reverse=True)]
def numpyDataTypeToC(dtype):
"""Mapping numpy data types to C data types"""
if dtype == np.float64:
return "double"
elif dtype == np.float32:
return "float"
elif dtype == np.int32:
return "int"
raise NotImplementedError()
def offsetComponentToDirectionString(coordinateId, value):
"""
Translates numerical offset to string notation.
x offsets are labeled with east 'E' and 'W',
y offsets with north 'N' and 'S' and
z offsets with top 'T' and bottom 'B'
If the absolute value of the offset is bigger than 1, this number is prefixed.
:param coordinateId: integer 0, 1 or 2 standing for x,y and z
:param value: integer offset
Example:
>>> offsetComponentToDirectionString(0, 1)
'E'
>>> offsetComponentToDirectionString(1, 2)
'2N'
"""
nameComponents = (('W', 'E'), # west, east
('S', 'N'), # south, north
('B', 'T'), # bottom, top
)
if value == 0:
result = ""
elif value < 0:
result = nameComponents[coordinateId][0]
else:
result = nameComponents[coordinateId][1]
if abs(value) > 1:
result = "%d%s" % (abs(value), result)
return result
def offsetToDirectionString(offsetTuple):
"""
Translates numerical offset to string notation.
For details see :func:`offsetComponentToDirectionString`
:param offsetTuple: 3-tuple with x,y,z offset
Example:
>>> offsetToDirectionString([1, -1, 0])
'SE'
>>> offsetToDirectionString(([-3, 0, -2]))
'2B3W'
"""
names = ["", "", ""]
for i in range(len(offsetTuple)):
names[i] = offsetComponentToDirectionString(i, offsetTuple[i])
name = "".join(reversed(names))
if name == "":
name = "C"
return name
def directionStringToOffset(directionStr, dim=3):
"""
Reverse mapping of :func:`offsetToDirectionString`
:param directionStr: string representation of offset
:param dim: dimension of offset, i.e the length of the returned list
>>> directionStringToOffset('NW', dim=3)
array([-1, 1, 0])
>>> directionStringToOffset('NW', dim=2)
array([-1, 1])
>>> directionStringToOffset(offsetToDirectionString([3,-2,1]))
array([ 3, -2, 1])
"""
offsetMap = {
'C': np.array([0, 0, 0]),
'W': np.array([-1, 0, 0]),
'E': np.array([1, 0, 0]),
'S': np.array([0, -1, 0]),
'N': np.array([0, 1, 0]),
'B': np.array([0, 0, -1]),
'T': np.array([0, 0, 1]),
}
offset = np.array([0, 0, 0])
while len(directionStr) > 0:
factor = 1
firstNonDigit = 0
while directionStr[firstNonDigit].isdigit():
firstNonDigit += 1
if firstNonDigit > 0:
factor = int(directionStr[:firstNonDigit])
directionStr = directionStr[firstNonDigit:]
curOffset = offsetMap[directionStr[0]]
offset += factor * curOffset
directionStr = directionStr[1:]
return offset[:dim]
import numpy as np
import sympy as sp
from pystencils.field import Field
def __upDownOffsets(d, dim):
coord = [0] * dim
coord[d] = 1
up = np.array(coord, dtype=np.int)
coord[d] = -1
down = np.array(coord, dtype=np.int)
return up, down
def grad(var, dim=3):
"""Gradients are represented as a special symbol:
e.g. :math:`\nabla x = (x^\Delta^0, x^\Delta^1, x^\Delta^2)`
r"""
Gradients are represented as a special symbol:
e.g. :math:`\nabla x = (x^{\Delta 0}, x^{\Delta 1}, x^{\Delta 2})`
This function takes a symbol and creates the gradient symbols according to convention above
:param var: symbol to take the gradient of
:param dim: dimension (length) of the gradient vector
"""
......@@ -29,7 +22,8 @@ def grad(var, dim=3):
def discretizeCenter(term, symbolsToFieldDict, dx, dim=3):
"""
Expects term that contains given symbols and gradient components of these symbols and replaces them
by field accesses. Gradients are replaced centralized approximations: (upper neighbor - lower neighbor ) / ( 2*dx).
by field accesses. Gradients are replaced centralized approximations:
``(upper neighbor - lower neighbor ) / ( 2*dx)``
:param term: term where symbols and gradient(symbol) should be replaced
:param symbolsToFieldDict: mapping of symbols to Field
:param dx: width and height of one cell
......@@ -57,15 +51,6 @@ def discretizeCenter(term, symbolsToFieldDict, dx, dim=3):
return term.subs(substitutions)
def fastSubs(term, subsDict):
def visit(expr):
if expr in subsDict:
return subsDict[expr]
paramList = [visit(a) for a in expr.args]
return expr if not paramList else expr.func(*paramList)
return visit(term)
def discretizeStaggered(term, symbolsToFieldDict, coordinate, coordinateOffset, dx, dim=3):
"""
Expects term that contains given symbols and gradient components of these symbols and replaces them
......@@ -122,13 +107,14 @@ def discretizeDivergence(vectorTerm, symbolsToFieldDict, dx):
Computes discrete divergence of symbolic vector
:param vectorTerm: sequence of terms, interpreted as vector
:param symbolsToFieldDict: mapping of symbols to Field
:param dx: length of a cell
Example: Laplace stencil
>>> x, dx = sp.symbols("x dx")
>>> gradx = grad(x, dim=3)
>>> f = Field.createGeneric('f', spatialDimensions=3)
>>> sp.simplify(discretizeDivergence(gradx, x, f, dx))
(f_B - 6*f_C + f_E + f_N + f_S + f_T + f_W)/dx
>>> x, dx = sp.symbols("x dx")
>>> gradX = grad(x, dim=3)
>>> f = Field.createGeneric('f', spatialDimensions=3)
>>> sp.simplify(discretizeDivergence(gradX, x, f, dx))
(f_B - 6*f_C + f_E + f_N + f_S + f_T + f_W)/dx
"""
dim = len(vectorTerm)
result = 0
......@@ -136,3 +122,23 @@ def discretizeDivergence(vectorTerm, symbolsToFieldDict, dx):
for offset in [-1, 1]:
result += offset * discretizeStaggered(vectorTerm[d], symbolsToFieldDict, d, offset, dx, dim)
return result
def __upDownOffsets(d, dim):
coord = [0] * dim
coord[d] = 1
up = np.array(coord, dtype=np.int)
coord[d] = -1
down = np.array(coord, dtype=np.int)
return up, down
def fastSubs(term, subsDict):
"""Similar to sympy subs function.
This version is much faster for big substitution dictionaries than sympy version"""
def visit(expr):
if expr in subsDict:
return subsDict[expr]
paramList = [visit(a) for a in expr.args]
return expr if not paramList else expr.func(*paramList)
return visit(term)
......@@ -7,14 +7,12 @@ from pystencils.typedsymbol import TypedSymbol
import pystencils.ast as ast
# --------------------------------------- Factory Functions ------------------------------------------------------------
def makeLoopOverDomain(body, functionName):
"""
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
:return: LoopOverCoordinate instance with nested loops, ordered according to field layouts
:return: :class:`LoopOverCoordinate` instance with nested loops, ordered according to field layouts
"""
# find correct ordering by inspecting participating FieldAccesses
fieldAccesses = body.atoms(Field.Access)
......@@ -45,9 +43,27 @@ def makeLoopOverDomain(body, functionName):
return ast.KernelFunction(currentBody, functionName)
# --------------------------------------- Transformations --------------------------------------------------------------
def createIntermediateBasePointer(fieldAccess, coordinates, previousPtr):
r"""
Addressing elements in structured arrays are done with :math:`ptr\left[ \sum_i c_i \cdot s_i \right]`
where :math:`c_i` is the coordinate value and :math:`s_i` the stride of a coordinate.
The sum can be split up into multiple parts, such that parts of it can be pulled before loops.
This function creates such an access for coordinates :math:`i \in \mbox{coordinates}`.
Returns a new typed symbol, where the name encodes which coordinates have been resolved.
:param fieldAccess: instance of :class:`pystencils.field.Field.Access` which provides strides and offsets
:param coordinates: mapping of coordinate ids to its value, where stride*value is calculated
:param previousPtr: the pointer which is dereferenced
:return: tuple with the new pointer symbol and the calculated offset
Example:
>>> field = Field.createGeneric('myfield', spatialDimensions=2, indexDimensions=1)
>>> x, y = sp.symbols("x y")
>>> prevPointer = TypedSymbol("ptr", "double")
>>> createIntermediateBasePointer(field[1,-2](5), {0: x}, prevPointer)
(ptr_E, x*fstride_myfield[0] + fstride_myfield[0])
>>> createIntermediateBasePointer(field[1,-2](5), {0: x, 1 : y }, prevPointer)
(ptr_E_2S, x*fstride_myfield[0] + y*fstride_myfield[1] + fstride_myfield[0] - 2*fstride_myfield[1])
"""
field = fieldAccess.field
offset = 0
......@@ -79,15 +95,24 @@ def createIntermediateBasePointer(fieldAccess, coordinates, previousPtr):
def parseBasePointerInfo(basePointerSpecification, loopOrder, field):
"""
Creates base pointer specification for :func:`resolveFieldAccesses` function.
Specification of how many and which intermediate pointers are created for a field access.
For example [ (0), (2,3,)] creates on base pointer for coordinates 2 and 3 and writes the offset for coordinate
zero directly in the field access. These specifications are more sensible defined dependent on the loop ordering.
This function translates more readable version into the specification above.
Allowed specifications:
"spatialInner<int>" spatialInner0 is the innermost loop coordinate, spatialInner1 the loop enclosing the innermost