serial_datahandling.py 13.9 KB
Newer Older
Martin Bauer's avatar
Martin Bauer committed
1
import itertools
2
3
4

import numpy as np

5
from pystencils import Field
6
from pystencils.datahandling.datahandling_interface import DataHandling
7
from pystencils.field import layoutStringToTuple, spatialLayoutStringToTuple, createNumpyArrayWithLayout
Martin Bauer's avatar
Martin Bauer committed
8
from pystencils.parallel.blockiteration import SerialBlock
Martin Bauer's avatar
Martin Bauer committed
9
10
11
12
13
from pystencils.slicing import normalizeSlice, removeGhostLayers
from pystencils.utils import DotDict

try:
    import pycuda.gpuarray as gpuarray
14
    import pycuda.autoinit
Martin Bauer's avatar
Martin Bauer committed
15
16
17
except ImportError:
    gpuarray = None

18

Martin Bauer's avatar
Martin Bauer committed
19
20
class SerialDataHandling(DataHandling):

21
    def __init__(self, domainSize, defaultGhostLayers=1, defaultLayout='SoA', periodicity=False):
Martin Bauer's avatar
Martin Bauer committed
22
23
24
25
26
27
28
        """
        Creates a data handling for single node simulations

        :param domainSize: size of the spatial domain as tuple
        :param defaultGhostLayers: nr of ghost layers used if not specified in add() method
        :param defaultLayout: layout used if no layout is given to add() method
        """
29
        super(SerialDataHandling, self).__init__()
Martin Bauer's avatar
Martin Bauer committed
30
31
32
33
34
35
        self._domainSize = tuple(domainSize)
        self.defaultGhostLayers = defaultGhostLayers
        self.defaultLayout = defaultLayout
        self._fields = DotDict()
        self.cpuArrays = DotDict()
        self.gpuArrays = DotDict()
36
37
38
39
        self.customDataCpu = DotDict()
        self.customDataGpu = DotDict()
        self._customDataTransferFunctions = {}

40
41
42
43
44
45
        if periodicity is None or periodicity is False:
            periodicity = [False] * self.dim
        if periodicity is True:
            periodicity = [True] * self.dim

        self._periodicity = periodicity
Martin Bauer's avatar
Martin Bauer committed
46
47
48
49
50
51
        self._fieldInformation = {}

    @property
    def dim(self):
        return len(self._domainSize)

Martin Bauer's avatar
Martin Bauer committed
52
53
54
55
56
57
58
59
    @property
    def shape(self):
        return self._domainSize

    @property
    def periodicity(self):
        return self._periodicity

Martin Bauer's avatar
Martin Bauer committed
60
61
62
63
    @property
    def fields(self):
        return self._fields

64
65
66
    def ghostLayersOfField(self, name):
        return self._fieldInformation[name]['ghostLayers']

67
68
69
    def fSize(self, name):
        return self._fieldInformation[name]['fSize']

70
71
    def addArray(self, name, fSize=1, dtype=np.float64, latexName=None, ghostLayers=None, layout=None,
                 cpu=True, gpu=False):
Martin Bauer's avatar
Martin Bauer committed
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
        if ghostLayers is None:
            ghostLayers = self.defaultGhostLayers
        if layout is None:
            layout = self.defaultLayout

        kwargs = {
            'shape': tuple(s + 2 * ghostLayers for s in self._domainSize),
            'dtype': dtype,
        }
        self._fieldInformation[name] = {
            'ghostLayers': ghostLayers,
            'fSize': fSize,
            'layout': layout,
            'dtype': dtype,
        }

        if fSize > 1:
            kwargs['shape'] = kwargs['shape'] + (fSize,)
            indexDimensions = 1
91
            layoutTuple = layoutStringToTuple(layout, self.dim+1)
Martin Bauer's avatar
Martin Bauer committed
92
93
        else:
            indexDimensions = 0
94
95
96
97
            layoutTuple = spatialLayoutStringToTuple(layout, self.dim)

        # cpuArr is always created - since there is no createPycudaArrayWithLayout()
        cpuArr = createNumpyArrayWithLayout(layout=layoutTuple, **kwargs)
98
99
        cpuArr.fill(np.inf)

Martin Bauer's avatar
Martin Bauer committed
100
101
102
        if cpu:
            if name in self.cpuArrays:
                raise ValueError("CPU Field with this name already exists")
103
            self.cpuArrays[name] = cpuArr
Martin Bauer's avatar
Martin Bauer committed
104
105
106
        if gpu:
            if name in self.gpuArrays:
                raise ValueError("GPU Field with this name already exists")
107
            self.gpuArrays[name] = gpuarray.to_gpu(cpuArr)
Martin Bauer's avatar
Martin Bauer committed
108

Martin Bauer's avatar
Martin Bauer committed
109
110
        assert all(f.name != name for f in self.fields.values()), "Symbolic field with this name already exists"
        self.fields[name] = Field.createFixedSize(name, shape=kwargs['shape'], indexDimensions=indexDimensions,
111
                                                  dtype=kwargs['dtype'], layout=layoutTuple)
Martin Bauer's avatar
Martin Bauer committed
112
        self.fields[name].latexName = latexName
113
        return self.fields[name]
Martin Bauer's avatar
Martin Bauer committed
114

115
116
    def addCustomData(self, name, cpuCreationFunction,
                      gpuCreationFunction=None, cpuToGpuTransferFunc=None, gpuToCpuTransferFunc=None):
117
118

        if cpuCreationFunction and gpuCreationFunction:
119
120
121
122
            if cpuToGpuTransferFunc is None or gpuToCpuTransferFunc is None:
                raise ValueError("For GPU data, both transfer functions have to be specified")
            self._customDataTransferFunctions[name] = (cpuToGpuTransferFunc, gpuToCpuTransferFunc)

123
124
125
126
127
128
129
130
131
        assert name not in self.customDataCpu
        if cpuCreationFunction:
            assert name not in self.cpuArrays
            self.customDataCpu[name] = cpuCreationFunction()

        if gpuCreationFunction:
            assert name not in self.gpuArrays
            self.customDataGpu[name] = gpuCreationFunction()

Martin Bauer's avatar
Martin Bauer committed
132
133
    def hasData(self, name):
        return name in self.fields
134

135
    def addArrayLike(self, name, nameOfTemplateField, latexName=None, cpu=True, gpu=False):
Martin Bauer's avatar
Martin Bauer committed
136
        return self.addArray(name, latexName=latexName, cpu=cpu, gpu=gpu, **self._fieldInformation[nameOfTemplateField])
Martin Bauer's avatar
Martin Bauer committed
137

Martin Bauer's avatar
Martin Bauer committed
138
    def iterate(self, sliceObj=None, gpu=False, ghostLayers=True, innerGhostLayers=True):
139
        if ghostLayers is True:
140
            ghostLayers = self.defaultGhostLayers
141
142
        elif ghostLayers is False:
            ghostLayers = 0
143
144
        elif isinstance(ghostLayers, str):
            ghostLayers = self.ghostLayersOfField(ghostLayers)
145

Martin Bauer's avatar
Martin Bauer committed
146
        if sliceObj is None:
147
148
149
150
151
            sliceObj = (slice(None, None, None),) * self.dim
        sliceObj = normalizeSlice(sliceObj, tuple(s + 2 * ghostLayers for s in self._domainSize))
        sliceObj = tuple(s if type(s) is slice else slice(s, s+1, None) for s in sliceObj)

        arrays = self.gpuArrays if gpu else self.cpuArrays
152
153
        customDataDict = self.customDataGpu if gpu else self.customDataCpu
        iterDict = customDataDict.copy()
154
155
156
157
158
159
160
161
162
        for name, arr in arrays.items():
            fieldGls = self._fieldInformation[name]['ghostLayers']
            if fieldGls < ghostLayers:
                continue
            arr = removeGhostLayers(arr, indexDimensions=len(arr.shape) - self.dim, ghostLayers=fieldGls-ghostLayers)
            iterDict[name] = arr

        offset = tuple(s.start - ghostLayers for s in sliceObj)
        yield SerialBlock(iterDict, offset, sliceObj)
163

164
165
166
167
168
169
    def gatherArray(self, name, sliceObj=None, ghostLayers=False, **kwargs):
        glToRemove = self._fieldInformation[name]['ghostLayers']
        if isinstance(ghostLayers, int):
            glToRemove -= ghostLayers
        if ghostLayers is True:
            glToRemove = 0
170
        arr = self.cpuArrays[name]
171
        indDimensions = self.fields[name].indexDimensions
Martin Bauer's avatar
Martin Bauer committed
172
173
        spatialDimensions = self.fields[name].spatialDimensions

174
        arr = removeGhostLayers(arr, indexDimensions=indDimensions, ghostLayers=glToRemove)
175

176
        if sliceObj is not None:
Martin Bauer's avatar
Martin Bauer committed
177
178
179
180
            normalizedSlice = normalizeSlice(sliceObj[:spatialDimensions], arr.shape[:spatialDimensions])
            normalizedSlice = tuple(s if type(s) is slice else slice(s, s + 1, None) for s in normalizedSlice)
            normalizedSlice += sliceObj[spatialDimensions:]
            arr = arr[normalizedSlice]
181
182
183
184
        else:
            arr = arr.view()
        arr.flags.writeable = False
        return arr
Martin Bauer's avatar
Martin Bauer committed
185
186
187
188
189
190
191
192

    def swap(self, name1, name2, gpu=False):
        if not gpu:
            self.cpuArrays[name1], self.cpuArrays[name2] = self.cpuArrays[name2], self.cpuArrays[name1]
        else:
            self.gpuArrays[name1], self.gpuArrays[name2] = self.gpuArrays[name2], self.gpuArrays[name1]

    def allToCpu(self):
193
        for name in (self.cpuArrays.keys() & self.gpuArrays.keys()) | self._customDataTransferFunctions.keys():
Martin Bauer's avatar
Martin Bauer committed
194
195
196
            self.toCpu(name)

    def allToGpu(self):
197
        for name in (self.cpuArrays.keys() & self.gpuArrays.keys()) | self._customDataTransferFunctions.keys():
Martin Bauer's avatar
Martin Bauer committed
198
199
            self.toGpu(name)

200
    def runKernel(self, kernelFunc, *args, **kwargs):
Martin Bauer's avatar
Martin Bauer committed
201
        dataUsedInKernel = [p.fieldName
202
                            for p in kernelFunc.parameters if p.isFieldPtrArgument and p.fieldName not in kwargs]
203
204
205
206
207
        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)

Martin Bauer's avatar
Martin Bauer committed
208
    def toCpu(self, name):
209
210
211
212
213
        if name in self._customDataTransferFunctions:
            transferFunc = self._customDataTransferFunctions[name][1]
            transferFunc(self.customDataGpu[name], self.customDataCpu[name])
        else:
            self.gpuArrays[name].get(self.cpuArrays[name])
Martin Bauer's avatar
Martin Bauer committed
214
215

    def toGpu(self, name):
216
217
218
219
220
        if name in self._customDataTransferFunctions:
            transferFunc = self._customDataTransferFunctions[name][0]
            transferFunc(self.customDataGpu[name], self.customDataCpu[name])
        else:
            self.gpuArrays[name].set(self.cpuArrays[name])
221

222
223
224
    def isOnGpu(self, name):
        return name in self.gpuArrays

225
    def synchronizationFunctionCPU(self, names, stencilName=None, **kwargs):
226
        return self.synchronizationFunction(names, stencilName, 'cpu')
227
228

    def synchronizationFunctionGPU(self, names, stencilName=None, **kwargs):
229
        return self.synchronizationFunction(names, stencilName, 'gpu')
230

231
    def synchronizationFunction(self, names, stencil=None, target='cpu'):
232
233
234
235
236
        assert target in ('cpu', 'gpu')
        if not hasattr(names, '__len__') or type(names) is str:
            names = [names]

        filteredStencil = []
Martin Bauer's avatar
Martin Bauer committed
237
238
        neighbors = [-1, 0, 1]

239
        if (stencil is None and self.dim == 2) or (stencil is not None and stencil.startswith('D2')):
Martin Bauer's avatar
Martin Bauer committed
240
            directions = itertools.product(*[neighbors] * 2)
241
        elif (stencil is None and self.dim == 3) or (stencil is not None and stencil.startswith('D3')):
Martin Bauer's avatar
Martin Bauer committed
242
243
244
245
246
            directions = itertools.product(*[neighbors] * 3)
        else:
            raise ValueError("Invalid stencil")

        for direction in directions:
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
            useDirection = True
            if direction == (0, 0) or direction == (0, 0, 0):
                useDirection = False
            for component, periodicity in zip(direction, self._periodicity):
                if not periodicity and component != 0:
                    useDirection = False
            if useDirection:
                filteredStencil.append(direction)

        resultFunctors = []
        for name in names:
            gls = self._fieldInformation[name]['ghostLayers']
            if len(filteredStencil) > 0:
                if target == 'cpu':
                    from pystencils.slicing import getPeriodicBoundaryFunctor
                    resultFunctors.append(getPeriodicBoundaryFunctor(filteredStencil, ghostLayers=gls))
                else:
                    from pystencils.gpucuda.periodicity import getPeriodicBoundaryFunctor
                    resultFunctors.append(getPeriodicBoundaryFunctor(filteredStencil, self._domainSize,
                                                                     indexDimensions=self.fields[name].indexDimensions,
                                                                     indexDimShape=self._fieldInformation[name]['fSize'],
                                                                     dtype=self.fields[name].dtype.numpyDtype,
                                                                     ghostLayers=gls))

        if target == 'cpu':
            def resultFunctor():
Martin Bauer's avatar
Martin Bauer committed
273
                for name, func in zip(names, resultFunctors):
274
275
276
                    func(pdfs=self.cpuArrays[name])
        else:
            def resultFunctor():
Martin Bauer's avatar
Martin Bauer committed
277
                for name, func in zip(names, resultFunctors):
278
279
280
                    func(pdfs=self.gpuArrays[name])

        return resultFunctor
Martin Bauer's avatar
Martin Bauer committed
281

282
283
284
285
286
287
288
289
    @property
    def arrayNames(self):
        return tuple(self.fields.keys())

    @property
    def customDataNames(self):
        return tuple(self.customDataCpu.keys())

Martin Bauer's avatar
Martin Bauer committed
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
    @staticmethod
    def reduceFloatSequence(sequence, operation, allReduce=False):
        return np.array(sequence)

    @staticmethod
    def reduceIntSequence(sequence):
        return np.array(sequence)

    def vtkWriter(self, fileName, dataNames, ghostLayers=False):
        from pystencils.vtk import imageToVTK

        def writer(step):
            fullFileName = "%s_%08d" % (fileName, step,)
            cellData = {}
            for name in dataNames:
                field = self._getFieldWithGivenNumberOfGhostLayers(name, ghostLayers)
                if self.dim == 2:
                    field = field[:, :, np.newaxis]
                if len(field.shape) == 3:
                    field = np.ascontiguousarray(field)
                elif len(field.shape) == 4:
                    field = [np.ascontiguousarray(field[..., i]) for i in range(field.shape[-1])]
                    if len(field) == 2:
                        field.append(np.zeros_like(field[0]))
                    field = tuple(field)
                else:
                    assert False
                cellData[name] = field
            imageToVTK(fullFileName, cellData=cellData)
        return writer

    def vtkWriterFlags(self, fileName, dataName, masksToName, ghostLayers=False):
        from pystencils.vtk import imageToVTK

        def writer(step):
            fullFileName = "%s_%08d" % (fileName, step,)
            field = self._getFieldWithGivenNumberOfGhostLayers(dataName, ghostLayers)
            if self.dim == 2:
                field = field[:, :, np.newaxis]
            cellData = {name: np.ascontiguousarray(np.bitwise_and(field, mask) > 0, dtype=np.uint8)
                        for mask, name in masksToName.items()}
            imageToVTK(fullFileName, cellData=cellData)

        return writer

    def _getFieldWithGivenNumberOfGhostLayers(self, name, ghostLayers):
        actualGhostLayers = self.ghostLayersOfField(name)
        if ghostLayers is True:
            ghostLayers = actualGhostLayers

        glToRemove = actualGhostLayers - ghostLayers
        indDims = 1 if self._fieldInformation[name]['fSize'] > 1 else 0
        return removeGhostLayers(self.cpuArrays[name], indDims, glToRemove)

344