diff --git a/cpu/kernelcreation.py b/cpu/kernelcreation.py index 6407382bae6f7ef64af77a354f009eccc8f8c6bb..c959afe09a33313cac4f52a3edc5bbcfd915c5c5 100644 --- a/cpu/kernelcreation.py +++ b/cpu/kernelcreation.py @@ -6,7 +6,8 @@ from pystencils.field import Field import pystencils.ast as ast -def createKernel(listOfEquations, functionName="kernel", typeForSymbol=None, splitGroups=(), iterationSlice=None): +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. @@ -21,6 +22,10 @@ def createKernel(listOfEquations, functionName="kernel", typeForSymbol=None, spl :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: @@ -43,7 +48,7 @@ def createKernel(listOfEquations, functionName="kernel", typeForSymbol=None, spl field.setReadOnly() body = ast.Block(assignments) - code = makeLoopOverDomain(body, functionName, iterationSlice=iterationSlice) + code = makeLoopOverDomain(body, functionName, iterationSlice=iterationSlice, ghostLayers=ghostLayers) if splitGroups: typedSplitGroups = [[typeSymbol(s) for s in splitGroup] for splitGroup in splitGroups] diff --git a/field.py b/field.py index 3507cd02e91ae84ff1cc1d96c65446b5e48a0acf..9c5dfba03384a08cc08ab46a44e370734ec00894 100644 --- a/field.py +++ b/field.py @@ -156,6 +156,11 @@ class Field: def __repr__(self): return self._fieldName + def neighbor(self, coordId, offset): + offsetList = [0] * self.spatialDimensions + offsetList[coordId] = offset + return Field.Access(self, tuple(offsetList)) + def __getitem__(self, offset): if type(offset) is np.ndarray: offset = tuple(offset) diff --git a/transformations.py b/transformations.py index 1443908037d4b025d201bdd3893e487a11c50032..634d577607f9d952e828f609dec7a7a4d48c86e9 100644 --- a/transformations.py +++ b/transformations.py @@ -9,12 +9,15 @@ from pystencils.slicing import normalizeSlice import pystencils.ast as ast -def makeLoopOverDomain(body, functionName, iterationSlice=None): +def makeLoopOverDomain(body, functionName, iterationSlice=None, ghostLayers=None): """ 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 :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:`LoopOverCoordinate` instance with nested loops, ordered according to field layouts """ # find correct ordering by inspecting participating FieldAccesses @@ -23,9 +26,6 @@ def makeLoopOverDomain(body, functionName, iterationSlice=None): fields = set(fieldList) loopOrder = getOptimalLoopOrdering(fields) - # find number of required ghost layers - requiredGhostLayers = max([fa.requiredGhostLayers for fa in fieldAccesses]) - shapes = set([f.spatialShape for f in fields]) if len(shapes) > 1: @@ -39,12 +39,16 @@ def makeLoopOverDomain(body, functionName, iterationSlice=None): if iterationSlice is not None: iterationSlice = normalizeSlice(iterationSlice, shape) + if ghostLayers is None: + requiredGhostLayers = max([fa.requiredGhostLayers for fa in fieldAccesses]) + ghostLayers = [(requiredGhostLayers, requiredGhostLayers)] * len(loopOrder) + currentBody = body lastLoop = None for i, loopCoordinate in enumerate(loopOrder): if iterationSlice is None: - begin = requiredGhostLayers - end = shape[loopCoordinate] - requiredGhostLayers + begin = ghostLayers[loopCoordinate][0] + end = shape[loopCoordinate] - ghostLayers[loopCoordinate][1] newLoop = ast.LoopOverCoordinate(currentBody, loopCoordinate, begin, end, 1) lastLoop = newLoop currentBody = ast.Block([lastLoop])