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

Run kernels directly on data handlings

parent 04784686
......@@ -166,7 +166,7 @@ class KernelFunction(Node):
self.ghostLayers = ghostLayers
# these variables are assumed to be global, so no automatic parameter is generated for them
self.globalVariables = set()
self.backend = ""
self.backend = backend
@property
def symbolsDefined(self):
......
import numpy as np
from abc import ABC, abstractmethod, abstractproperty
from abc import ABC, abstractmethod
from collections import defaultdict
from contextlib import contextmanager
from lbmpy.stencils import getStencil
from pystencils import Field, makeSlice
from pystencils import Field
from pystencils.parallel.blockiteration import BlockIterationInfo
from pystencils.slicing import normalizeSlice, removeGhostLayers
from pystencils.utils import DotDict
......@@ -122,11 +122,11 @@ class DataHandling(ABC):
:return: generator expression yielding the gathered field, the gathered field does not include any ghost layers
"""
def registerPreAccessFunction(self, name, function):
self._preAccessFunctions[name].append(function)
def registerPreAccessFunction(self, name, fct):
self._preAccessFunctions[name].append(fct)
def registerPostAccessFunction(self, name, function):
self._postAccessFunctions[name].append(function)
def registerPostAccessFunction(self, name, fct):
self._postAccessFunctions[name].append(fct)
@contextmanager
def accessWrapper(self, name):
......@@ -200,6 +200,7 @@ class SerialDataHandling(DataHandling):
self.defaultGhostLayers = defaultGhostLayers
self.defaultLayout = defaultLayout
self._fields = DotDict()
self._fieldLatexNameToDataName = {}
self.cpuArrays = DotDict()
self.gpuArrays = DotDict()
self.customDataCpu = DotDict()
......@@ -222,7 +223,8 @@ class SerialDataHandling(DataHandling):
def fields(self):
return self._fields
def addArray(self, name, fSize=1, dtype=np.float64, latexName=None, ghostLayers=None, layout=None, cpu=True, gpu=False):
def addArray(self, name, fSize=1, dtype=np.float64, latexName=None, ghostLayers=None, layout=None,
cpu=True, gpu=False):
if ghostLayers is None:
ghostLayers = self.defaultGhostLayers
if layout is None:
......@@ -262,6 +264,7 @@ class SerialDataHandling(DataHandling):
assert all(f.name != latexName for f in self.fields.values()), "Symbolic field with this name already exists"
self.fields[name] = Field.createFixedSize(latexName, shape=kwargs['shape'], indexDimensions=indexDimensions,
dtype=kwargs['dtype'], layout=kwargs['order'])
self._fieldLatexNameToDataName[latexName] = name
def addCustomData(self, name, cpuCreationFunction,
gpuCreationFunction=None, cpuToGpuTransferFunc=None, gpuToCpuTransferFunc=None):
......@@ -281,8 +284,10 @@ class SerialDataHandling(DataHandling):
self.addArray(name, latexName=latexName, cpu=cpu, gpu=gpu, **self._fieldInformation[nameOfTemplateField])
def accessArray(self, name, sliceObj=None, outerGhostLayers='all', **kwargs):
if outerGhostLayers == 'all':
if outerGhostLayers == 'all' or outerGhostLayers is True:
outerGhostLayers = self._fieldInformation[name]['ghostLayers']
elif outerGhostLayers is False:
outerGhostLayers = 0
if sliceObj is None:
sliceObj = [slice(None, None)] * self.dim
......@@ -296,7 +301,7 @@ class SerialDataHandling(DataHandling):
yield arr[sliceObj], BlockIterationInfo(None, tuple(s.start for s in sliceObj), sliceObj)
def accessCustomData(self, name):
yield self.customDataCpu[name], ((0,0,0)[:self.dim], self._domainSize)
yield self.customDataCpu[name], ((0, 0, 0)[:self.dim], self._domainSize)
def gatherArray(self, name, sliceObj=None, **kwargs):
with self.accessWrapper(name):
......@@ -322,6 +327,14 @@ class SerialDataHandling(DataHandling):
for name in (self.cpuArrays.keys() & self.gpuArrays.keys()) | self._customDataTransferFunctions.keys():
self.toGpu(name)
def runKernel(self, kernelFunc, *args, **kwargs):
dataUsedInKernel = [self._fieldLatexNameToDataName[p.fieldName]
for p in kernelFunc.ast.parameters if p.isFieldPtrArgument]
arrays = self.gpuArrays if kernelFunc.ast.backend == 'gpucuda' else self.cpuArrays
arrayParams = {name: arrays[name] for name in dataUsedInKernel}
arrayParams.update(kwargs)
kernelFunc(*args, **arrayParams)
def toCpu(self, name):
if name in self._customDataTransferFunctions:
transferFunc = self._customDataTransferFunctions[name][1]
......
......@@ -143,7 +143,7 @@ class Field(object):
return Field(fieldName, FieldType.GENERIC, npArray.dtype, spatialLayout, shape, strides)
@staticmethod
def createFixedSize(fieldName, shape, indexDimensions=0, dtype=np.float64, layout='numpy'):
def createFixedSize(fieldName, shape, indexDimensions=0, dtype=np.float64, layout='numpy', strides=None):
"""
Creates a field with fixed sizes i.e. can be called only with arrays of the same size and layout
......@@ -152,6 +152,7 @@ class Field(object):
:param indexDimensions: how many of the trailing dimensions are interpreted as index (as opposed to spatial)
:param dtype: numpy data type of the array the kernel is called with later
:param layout: full layout of array, not only spatial dimensions
:param strides: strides in bytes or None to automatically compute them from shape (assuming no padding)
"""
spatialDimensions = len(shape) - indexDimensions
assert spatialDimensions >= 1
......@@ -160,7 +161,11 @@ class Field(object):
layout = layoutStringToTuple(layout, spatialDimensions + indexDimensions)
shape = tuple(int(s) for s in shape)
strides = computeStrides(shape, layout)
if strides is None:
strides = computeStrides(shape, layout)
else:
assert len(strides) == len(shape)
strides = tuple([s // np.dtype(dtype).itemsize for s in strides])
npDataType = np.dtype(dtype)
if npDataType.fields is not None:
......@@ -244,7 +249,7 @@ class Field(object):
return Field.Access(self, tuple(offsetList))
def neighbors(self, stencil):
return [ self.__getitem__(dir) for s in stencil]
return [self.__getitem__(s) for s in stencil]
@property
def center(self):
......
......@@ -26,6 +26,8 @@ class ParallelDataHandling(DataHandling):
self.defaultGhostLayers = defaultGhostLayers
self.defaultLayout = defaultLayout
self._fields = DotDict() # maps name to symbolic pystencils field
self._fieldNameToCpuDataName = {}
self._fieldNameToGpuDataName = {}
self.dataNames = set()
self._dim = dim
self._fieldInformation = {}
......@@ -84,13 +86,18 @@ class ParallelDataHandling(DataHandling):
self._cpuGpuPairs.append((name, self.GPU_DATA_PREFIX + name))
blockBB = self.blocks.getBlockCellBB(self.blocks[0])
shape = tuple(s + 2 * ghostLayers for s in blockBB.size)
shape = tuple(s + 2 * ghostLayers for s in blockBB.size[:self.dim])
indexDimensions = 1 if fSize > 1 else 0
if indexDimensions == 1:
shape += (fSize, )
assert all(f.name != latexName for f in self.fields.values()), "Symbolic field with this name already exists"
self.fields[name] = Field.createFixedSize(latexName, shape, indexDimensions, dtype, layout)
self.fields[name] = Field.createGeneric(latexName, self.dim, dtype, indexDimensions, layout,
indexShape=(fSize,) if indexDimensions > 0 else None)
self._fieldNameToCpuDataName[latexName] = name
if gpu:
self._fieldNameToGpuDataName[latexName] = self.GPU_DATA_PREFIX + name
def hasData(self, name):
return name in self._fields
......@@ -108,17 +115,18 @@ class ParallelDataHandling(DataHandling):
def accessArray(self, name, sliceObj=None, innerGhostLayers='all', outerGhostLayers='all'):
fieldInfo = self._fieldInformation[name]
with self.accessWrapper(name):
if innerGhostLayers is 'all':
if innerGhostLayers is 'all' or innerGhostLayers is True:
innerGhostLayers = fieldInfo['ghostLayers']
if outerGhostLayers is 'all':
elif innerGhostLayers is False:
innerGhostLayers = 0
if outerGhostLayers is 'all' or innerGhostLayers is True:
outerGhostLayers = fieldInfo['ghostLayers']
elif outerGhostLayers is False:
outerGhostLayers = 0
glAccess = [innerGhostLayers, innerGhostLayers, 0 if self.dim == 2 else innerGhostLayers]
for iterInfo in slicedBlockIteration(self.blocks, sliceObj, innerGhostLayers, outerGhostLayers):
arr = wlb.field.toArray(iterInfo.block[name], withGhostLayers=innerGhostLayers)[iterInfo.localSlice]
if self.fields[name].indexDimensions == 0:
arr = arr[..., 0]
if self.dim == 2:
arr = arr[:, :, 0]
arr = wlb.field.toArray(iterInfo.block[name], withGhostLayers=glAccess)[iterInfo.localSlice]
arr = self._normalizeArrShape(arr, self.fields[name].indexDimensions)
yield arr, iterInfo
def accessCustomData(self, name):
......@@ -141,13 +149,30 @@ class ParallelDataHandling(DataHandling):
array = array[:, :, 0]
yield array
def _normalizeArrShape(self, arr, indexDimensions):
if indexDimensions == 0:
arr = arr[..., 0]
if self.dim == 2:
arr = arr[:, :, 0]
return arr
def runKernel(self, kernelFunc, *args, **kwargs):
fieldArguments = [p.fieldName for p in kernelFunc.ast.parameters if p.isFieldPtrArgument]
if kernelFunc.ast.backend == 'gpucuda':
nameMap = self._fieldNameToGpuDataName
toArray = wlb.cuda.toGpuArray
else:
nameMap = self._fieldNameToCpuDataName
toArray = wlb.field.toArray
dataUsedInKernel = [(nameMap[p.fieldName], self.fields[p.fieldName])
for p in kernelFunc.ast.parameters if p.isFieldPtrArgument]
for block in self.blocks:
fieldArgs = {fieldName: wlb.field.toArray(block[fieldName], withGhostLayers=True)
for fieldName in fieldArguments}
fieldArgs = {}
for dataName, f in dataUsedInKernel:
arr = toArray(block[dataName], withGhostLayers=[True, True, self.dim == 3])
arr = self._normalizeArrShape(arr, f.indexDimensions)
fieldArgs[f.name] = arr
fieldArgs.update(kwargs)
kernelFunc(*args, **kwargs)
kernelFunc(*args, **fieldArgs)
def toCpu(self, name):
if name in self._customDataTransferFunctions:
......
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