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

Attempt to generate kerncraft compatible code

parent c744ab1d
......@@ -334,6 +334,13 @@ class SympyAssignment(Node):
@property
def undefinedSymbols(self):
result = self.rhs.atoms(sp.Symbol)
# Add loop counters if there a field accesses
loopCounters = set()
for symbol in result:
if isinstance(symbol, Field.Access):
for i in range(len(symbol.offsets)):
loopCounters.add(LoopOverCoordinate.getLoopCounterSymbol(i))
result.update(loopCounters)
result.update(self._lhsSymbol.atoms(sp.Symbol))
return result
......
......@@ -52,9 +52,13 @@ class PrintNode(CustomCppCode):
class CBackend:
def __init__(self, cuda=False):
def __init__(self, cuda=False, sympyPrinter=None):
self.cuda = cuda
self.sympyPrinter = CustomSympyPrinter()
if sympyPrinter is None:
self.sympyPrinter = CustomSympyPrinter()
else:
self.sympyPrinter = sympyPrinter
self._indent = " "
def __call__(self, node):
......
from pystencils.transformations import makeLoopOverDomain, typingFromSympyInspection, \
typeAllEquations, moveConstantsBeforeLoop, getOptimalLoopOrdering
import pystencils.ast as ast
from pystencils.backends.cbackend import CBackend, CustomSympyPrinter
from pystencils import TypedSymbol
def createKerncraftCode(listOfEquations, typeForSymbol=None, ghostLayers=None):
"""
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 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 ghostLayers: a sequence of pairs for each coordinate with lower and upper nr of ghost layers
if None, the number of ghost layers is determined automatically and assumed to be equal for a
all dimensions
:return: :class:`pystencils.ast.KernelFunction` node
"""
if not typeForSymbol:
typeForSymbol = typingFromSympyInspection(listOfEquations, "double")
fieldsRead, fieldsWritten, assignments = typeAllEquations(listOfEquations, typeForSymbol)
allFields = fieldsRead.union(fieldsWritten)
optimalLoopOrder = getOptimalLoopOrdering(allFields)
cstyleLoopOrder = list(range(len(optimalLoopOrder)))
body = ast.Block(assignments)
code = makeLoopOverDomain(body, "kerncraft", ghostLayers=ghostLayers, loopOrder=cstyleLoopOrder)
moveConstantsBeforeLoop(code)
loopBody = code.body
printer = CBackend(sympyPrinter=ArraySympyPrinter())
FIXED_SIZES = ("XS", "YS", "ZS", "E1S", "E2S")
result = ""
for field in allFields:
sizesPermutation = [FIXED_SIZES[i] for i in field.layout]
suffix = "".join("[%s]" % (size,) for size in sizesPermutation)
result += "%s%s;\n" % (field.name, suffix)
# add parameter definitions
for s in loopBody.undefinedSymbols:
if isinstance(s, TypedSymbol):
result += "%s %s;\n" % (s.dtype, s.name)
for element in loopBody.args:
result += printer(element)
result += "\n"
return result
class ArraySympyPrinter(CustomSympyPrinter):
def _print_Access(self, fieldAccess):
""""""
Loop = ast.LoopOverCoordinate
coordinateValues = [Loop.getLoopCounterSymbol(i) + offset for i, offset in enumerate(fieldAccess.offsets)]
coordinateValues += list(fieldAccess.index)
permutedCoordinates = [coordinateValues[i] for i in fieldAccess.field.layout]
suffix = "".join("[%s]" % (self._print(a)) for a in permutedCoordinates)
return "%s%s" % (self._print(fieldAccess.field.name), suffix)
......@@ -45,14 +45,14 @@ def createKernel(listOfEquations, functionName="kernel", typeForSymbol=None, spl
readOnlyFields = set([f.name for f in fieldsRead - fieldsWritten])
body = ast.Block(assignments)
code = makeLoopOverDomain(body, functionName, iterationSlice=iterationSlice, ghostLayers=ghostLayers)
loopOrder = getOptimalLoopOrdering(allFields)
code = makeLoopOverDomain(body, functionName, iterationSlice=iterationSlice,
ghostLayers=ghostLayers, loopOrder=loopOrder)
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}
......
......@@ -44,7 +44,7 @@ def createCUDAKernel(listOfEquations, functionName="kernel", typeForSymbol=defau
coordMapping, getCallParameters = getLinewiseCoordinates(list(fieldsRead)[0], requiredGhostLayers)
allFields = fieldsRead.union(fieldsWritten)
basePointerInfo = [['spatialInner0']]
basePointerInfos = {f.name: parseBasePointerInfo(basePointerInfo, [0, 1, 2], f) for f in allFields}
basePointerInfos = {f.name: parseBasePointerInfo(basePointerInfo, [2, 1, 0], f) for f in allFields}
resolveFieldAccesses(code, readOnlyFields, fieldToFixedCoordinates={'src': coordMapping, 'dst': coordMapping},
fieldToBasePointerInfo=basePointerInfos)
# add the function which determines #blocks and #threads as additional member to KernelFunction node
......
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=(),
iterationSlice=None, ghostLayers=None):
"""
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`
:param iterationSlice: if not None, iteration is done only over this slice of the field
:param ghostLayers: a sequence of pairs for each coordinate with lower and upper nr of ghost layers
if None, the number of ghost layers is determined automatically and assumed to be equal for a
all dimensions
: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)
readOnlyFields = set([f.name for f in fieldsRead - fieldsWritten])
body = ast.Block(assignments)
loopOrder = getOptimalLoopOrdering(allFields)
code = makeLoopOverDomain(body, functionName, iterationSlice=iterationSlice,
ghostLayers=ghostLayers, loopOrder=loopOrder)
if splitGroups:
typedSplitGroups = [[typeSymbol(s) for s in splitGroup] for splitGroup in splitGroups]
splitInnerLoop(code, typedSplitGroups)
basePointerInfo = [['spatialInner0'], ['spatialInner1']]
basePointerInfos = {field.name: parseBasePointerInfo(basePointerInfo, loopOrder, field) for field in allFields}
resolveFieldAccesses(code, readOnlyFields, fieldToBasePointerInfo=basePointerInfos)
moveConstantsBeforeLoop(code)
return code
\ No newline at end of file
......@@ -9,7 +9,7 @@ from pystencils.slicing import normalizeSlice
import pystencils.ast as ast
def makeLoopOverDomain(body, functionName, iterationSlice=None, ghostLayers=None):
def makeLoopOverDomain(body, functionName, iterationSlice=None, ghostLayers=None, loopOrder=None):
"""
Uses :class:`pystencils.field.Field.Access` to create (multiple) loops around given AST.
:param body: list of nodes
......@@ -18,13 +18,16 @@ def makeLoopOverDomain(body, functionName, iterationSlice=None, ghostLayers=None
:param ghostLayers: a sequence of pairs for each coordinate with lower and upper nr of ghost layers
if None, the number of ghost layers is determined automatically and assumed to be equal for a
all dimensions
:param loopOrder: loop ordering from outer to inner loop (optimal ordering is same as layout)
:return: :class:`LoopOverCoordinate` instance with nested loops, ordered according to field layouts
"""
# find correct ordering by inspecting participating FieldAccesses
fieldAccesses = body.atoms(Field.Access)
fieldList = [e.field for e in fieldAccesses]
fields = set(fieldList)
loopOrder = getOptimalLoopOrdering(fields)
if loopOrder is None:
loopOrder = getOptimalLoopOrdering(fields)
shapes = set([f.spatialShape for f in fields])
......@@ -45,7 +48,7 @@ def makeLoopOverDomain(body, functionName, iterationSlice=None, ghostLayers=None
currentBody = body
lastLoop = None
for i, loopCoordinate in enumerate(loopOrder):
for i, loopCoordinate in enumerate(reversed(loopOrder)):
if iterationSlice is None:
begin = ghostLayers[loopCoordinate][0]
end = shape[loopCoordinate] - ghostLayers[loopCoordinate][1]
......@@ -133,12 +136,13 @@ def parseBasePointerInfo(basePointerSpecification, loopOrder, field):
- "<int>": specifying directly the coordinate
:param basePointerSpecification: nested list with above specifications
:param loopOrder: list with ordering of loops from inner to outer
:param loopOrder: list with ordering of loops from outer to inner
:param field:
:return: list of tuples that can be passed to :func:`resolveFieldAccesses`
"""
result = []
specifiedCoordinates = set()
loopOrder = list(reversed(loopOrder))
for specGroup in basePointerSpecification:
newGroup = []
......@@ -461,7 +465,7 @@ def getOptimalLoopOrdering(fields):
Determines the optimal loop order for a given set of fields.
If the fields have different memory layout or different sizes an exception is thrown.
:param fields: sequence of fields
:return: list of coordinate ids, where the first list entry should be the innermost loop
:return: list of coordinate ids, where the first list entry should be the outermost loop
"""
assert len(fields) > 0
refField = next(iter(fields))
......@@ -473,7 +477,7 @@ def getOptimalLoopOrdering(fields):
if len(layouts) > 1:
raise ValueError("Due to different layout of the fields no optimal loop ordering exists " + str(layouts))
layout = list(layouts)[0]
return list(reversed(layout))
return list(layout)
def getLoopHierarchy(astNode):
......@@ -487,4 +491,4 @@ def getLoopHierarchy(astNode):
node = getNextParentOfType(node, ast.LoopOverCoordinate)
if node:
result.append(node.coordinateToLoopOver)
return result
\ No newline at end of file
return reversed(result)
\ No newline at end of file
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