diff --git a/astnodes.py b/astnodes.py index 8210771c6d39afe4ea760383a615539d04506526..ed0771778c2b5acbf36dd97e479ace4cf8763b79 100644 --- a/astnodes.py +++ b/astnodes.py @@ -1,7 +1,7 @@ import sympy as sp from sympy.tensor import IndexedBase from pystencils.field import Field -from pystencils.types import TypedSymbol, createType, get_type_from_sympy +from pystencils.types import TypedSymbol, createType, get_type_from_sympy, createTypeFromString class Node(object): @@ -301,6 +301,16 @@ class LoopOverCoordinate(Node): def loopCounterName(self): return LoopOverCoordinate.getLoopCounterName(self.coordinateToLoopOver) + @staticmethod + def isLoopCounterSymbol(symbol): + prefix = LoopOverCoordinate.LOOP_COUNTER_NAME_PREFIX + if not symbol.name.startswith(prefix): + return None + if symbol.dtype != createTypeFromString('int'): + return None + coordinate = int(symbol.name[len(prefix)+1:]) + return coordinate + @staticmethod def getLoopCounterSymbol(coordinateToLoopOver): return TypedSymbol(LoopOverCoordinate.getLoopCounterName(coordinateToLoopOver), 'int') diff --git a/gpucuda/kernelcreation.py b/gpucuda/kernelcreation.py index 3d825fa21ae46e40f13a8be6a4cb264ab9db324b..4aa87969d4878a5200da5d7a364683011d7a7d82 100644 --- a/gpucuda/kernelcreation.py +++ b/gpucuda/kernelcreation.py @@ -1,6 +1,6 @@ from pystencils.gpucuda.indexing import BlockIndexing from pystencils.transformations import resolveFieldAccesses, typeAllEquations, parseBasePointerInfo, getCommonShape -from pystencils.astnodes import Block, KernelFunction, SympyAssignment +from pystencils.astnodes import Block, KernelFunction, SympyAssignment, LoopOverCoordinate from pystencils.types import TypedSymbol, BasicType, StructType from pystencils import Field @@ -25,7 +25,7 @@ def createCUDAKernel(listOfEquations, functionName="kernel", typeForSymbol=None, iterationSlice = [] if isinstance(ghostLayers, int): for i in range(len(commonShape)): - iterationSlice.append(slice(ghostLayers[i], -ghostLayers[i] if ghostLayers[i] > 0 else None)) + iterationSlice.append(slice(ghostLayers, -ghostLayers if ghostLayers > 0 else None)) else: for i in range(len(commonShape)): iterationSlice.append(slice(ghostLayers[i][0], -ghostLayers[i][1] if ghostLayers[i][1] > 0 else None)) @@ -46,6 +46,14 @@ def createCUDAKernel(listOfEquations, functionName="kernel", typeForSymbol=None, fieldToBasePointerInfo=basePointerInfos) # add the function which determines #blocks and #threads as additional member to KernelFunction node # this is used by the jit + + # If loop counter symbols have been explicitly used in the update equations (e.g. for built in periodicity), + # they are defined here + undefinedLoopCounters = {LoopOverCoordinate.isLoopCounterSymbol(s): s for s in ast.body.undefinedSymbols + if LoopOverCoordinate.isLoopCounterSymbol(s) is not None} + for i, loopCounter in undefinedLoopCounters.items(): + ast.body.insertFront(SympyAssignment(loopCounter, indexing.coordinates[i])) + ast.indexing = indexing return ast