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 ...@@ -5,21 +5,29 @@ from pystencils.typedsymbol import TypedSymbol
class Node: class Node:
"""Base class for all AST nodes"""
def __init__(self, parent=None): def __init__(self, parent=None):
self.parent = parent self.parent = parent
def args(self): def args(self):
"""Returns all arguments/children of this node"""
return [] return []
@property @property
def symbolsDefined(self): def symbolsDefined(self):
"""Set of symbols which are defined in this node or its children"""
return set() return set()
@property @property
def symbolsRead(self): def symbolsRead(self):
"""Set of symbols which are accessed/read in this node or its children"""
return set() return set()
def atoms(self, argType): def atoms(self, argType):
"""
Returns a set of all children which are an instance of the given argType
"""
result = set() result = set()
for arg in self.args: for arg in self.args:
if isinstance(arg, argType): if isinstance(arg, argType):
......
...@@ -4,6 +4,9 @@ from pystencils.ast import Node ...@@ -4,6 +4,9 @@ from pystencils.ast import Node
def printCCode(astNode): def printCCode(astNode):
"""
Prints the abstract syntax tree as C function
"""
printer = CBackend(cuda=False) printer = CBackend(cuda=False)
return printer(astNode) 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): ...@@ -131,6 +131,18 @@ def makePythonFunctionIncompleteParams(kernelFunctionNode, argumentDict):
def makePythonFunction(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 # build up list of CType arguments
try: try:
args = buildCTypeArgumentList(kernelFunctionNode, argumentDict) args = buildCTypeArgumentList(kernelFunctionNode, argumentDict)
...@@ -140,10 +152,3 @@ def makePythonFunction(kernelFunctionNode, argumentDict={}): ...@@ -140,10 +152,3 @@ def makePythonFunction(kernelFunctionNode, argumentDict={}):
func = compileAndLoad(kernelFunctionNode)[kernelFunctionNode.functionName] func = compileAndLoad(kernelFunctionNode)[kernelFunctionNode.functionName]
func.restype = None func.restype = None
return lambda: func(*args) 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 ...@@ -6,147 +6,31 @@ from sympy.tensor import IndexedBase
from pystencils.typedsymbol import TypedSymbol 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: class Field:
""" """
With fields one can formulate stencil-like update rules on structured grids. 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. 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 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 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: Dimensions:
A field has spatial and index dimensions, where the spatial dimensions come first. 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 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, 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 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: Indexing:
When accessing (indexing) a field the result is a FieldAccess which is derived from sympy Symbol. 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 () 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: Example without index dimensions:
>>> a = np.zeros([10, 10]) >>> a = np.zeros([10, 10])
...@@ -162,6 +46,7 @@ class Field: ...@@ -162,6 +46,7 @@ class Field:
Eq(dst_C^0, src_C^0) Eq(dst_C^0, src_C^0)
Eq(dst_C^1, src_S^1) Eq(dst_C^1, src_S^1)
Eq(dst_C^2, src_N^2) Eq(dst_C^2, src_N^2)
""" """
@staticmethod @staticmethod
def createFromNumpyArray(fieldName, npArray, indexDimensions=0): def createFromNumpyArray(fieldName, npArray, indexDimensions=0):
...@@ -194,8 +79,8 @@ class Field: ...@@ -194,8 +79,8 @@ class Field:
:param spatialDimensions: see documentation of Field :param spatialDimensions: see documentation of Field
:param indexDimensions: 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 :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 the outer loop loops over dimension 2, the second outer over dimension 1, and the inner loop
over dimension 0 over dimension 0
""" """
if not layout: if not layout:
layout = tuple(reversed(range(spatialDimensions))) layout = tuple(reversed(range(spatialDimensions)))
...@@ -384,3 +269,137 @@ class Field: ...@@ -384,3 +269,137 @@ class Field:
superClassContents = list(super(Field.Access, self)._hashable_content()) superClassContents = list(super(Field.Access, self)._hashable_content())
t = tuple([*superClassContents, hash(self._field), self._index] + self._offsets) t = tuple([*superClassContents, hash(self._field), self._index] + self._offsets)
return t 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 numpy as np
import sympy as sp import sympy as sp
from pystencils.field import Field 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): def grad(var, dim=3):
"""Gradients are represented as a special symbol: r"""
e.g. :math:`\nabla x = (x^\Delta^0, x^\Delta^1, x^\Delta^2)` 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 This function takes a symbol and creates the gradient symbols according to convention above
:param var: symbol to take the gradient of :param var: symbol to take the gradient of
:param dim: dimension (length) of the gradient vector :param dim: dimension (length) of the gradient vector
""" """
...@@ -29,7 +22,8 @@ def grad(var, dim=3): ...@@ -29,7 +22,8 @@ def grad(var, dim=3):
def discretizeCenter(term, symbolsToFieldDict, dx, dim=3): def discretizeCenter(term, symbolsToFieldDict, dx, dim=3):
""" """
Expects term that contains given symbols and gradient components of these symbols and replaces them 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 term: term where symbols and gradient(symbol) should be replaced
:param symbolsToFieldDict: mapping of symbols to Field :param symbolsToFieldDict: mapping of symbols to Field
:param dx: width and height of one cell :param dx: width and height of one cell
...@@ -57,15 +51,6 @@ def discretizeCenter(term, symbolsToFieldDict, dx, dim=3): ...@@ -57,15 +51,6 @@ def discretizeCenter(term, symbolsToFieldDict, dx, dim=3):
return term.subs(substitutions) 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): def discretizeStaggered(term, symbolsToFieldDict, coordinate, coordinateOffset, dx, dim=3):
""" """
Expects term that contains given symbols and gradient components of these symbols and replaces them Expects term that contains given symbols and gradient components of these symbols and replaces them
...@@ -122,13 +107,14 @@ def discretizeDivergence(vectorTerm, symbolsToFieldDict, dx): ...@@ -122,13 +107,14 @@ def discretizeDivergence(vectorTerm, symbolsToFieldDict, dx):
Computes discrete divergence of symbolic vector Computes discrete divergence of symbolic vector
:param vectorTerm: sequence of terms, interpreted as vector :param vectorTerm: sequence of terms, interpreted as vector
:param symbolsToFieldDict: mapping of symbols to Field :param symbolsToFieldDict: mapping of symbols to Field
:param dx: length of a cell
Example: Laplace stencil Example: Laplace stencil
>>> x, dx = sp.symbols("x dx") >>> x, dx = sp.symbols("x dx")
>>> gradx = grad(x, dim=3) >>> gradX = grad(x, dim=3)
>>> f = Field.createGeneric('f', spatialDimensions=3) >>> f = Field.createGeneric('f', spatialDimensions=3)
>>> sp.simplify(discretizeDivergence(gradx, x, f, dx)) >>> sp.simplify(discretizeDivergence(gradX, x, f, dx))
(f_B - 6*f_C + f_E + f_N + f_S + f_T + f_W)/dx (f_B - 6*f_C + f_E + f_N + f_S + f_T + f_W)/dx
""" """
dim = len(vectorTerm) dim = len(vectorTerm)
result = 0 result = 0
...@@ -136,3 +122,23 @@ def discretizeDivergence(vectorTerm, symbolsToFieldDict, dx): ...@@ -136,3 +122,23 @@ def discretizeDivergence(vectorTerm, symbolsToFieldDict, dx):
for offset in [-1, 1]: for offset in [-1, 1]:
result += offset * discretizeStaggered(vectorTerm[d], symbolsToFieldDict, d, offset, dx, dim) result += offset * discretizeStaggered(vectorTerm[d], symbolsToFieldDict, d, offset, dx, dim)
return result return result
def __upDownOffsets(d, dim):
coord = [0] * dim
coord[d] = 1