Attempt to generate kerncraft compatible code

ast.py

@@ -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
 

backends/cbackend.py

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

cpu/kerncraft.py

+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)

cpu/kernelcreation.py

@@ -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}
 

gpucuda/kernelcreation.py

@@ -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

llvm/kernelcreation.py

+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

transformations.py

@@ -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