Commit 9fe16eda authored by Martin Bauer's avatar Martin Bauer
Browse files

Added numpy array alignment option

parent 59d0f541
import numpy as np
def aligned_empty(shape, byteAlignment=32, dtype=np.float64, byteOffset=0, order='C', alignInnerCoordinate=True):
"""
Creates an aligned empty numpy array
:param shape: size of the array
:param byteAlignment: alignment in bytes, for the start address of the array holds (a % byteAlignment) == 0
:param dtype: numpy data type
:param byteOffset: offset in bytes for position that should be aligned i.e. (a+byteOffset) % byteAlignment == 0
typically used to align first inner cell instead of ghost layer
:param order: storage linearization order
:param alignInnerCoordinate: if True, the start of the innermost coordinate lines are aligned as well
:return:
"""
if (not alignInnerCoordinate) or (not hasattr(shape, '__len__')):
N = np.prod(shape)
d = np.dtype(dtype)
tmp = np.empty(N * d.itemsize + byteAlignment, dtype=np.uint8)
address = tmp.__array_interface__['data'][0]
offset = (byteAlignment - (address + byteOffset) % byteAlignment) % byteAlignment
return tmp[offset:offset + N * d.itemsize].view(dtype=d).reshape(shape, order=order)
else:
if order == 'C':
ndim0 = shape[-1]
dim0 = -1
ndim1 = shape[-2]
else:
ndim0 = shape[0]
dim0 = 0
ndim1 = shape[1]
d = np.dtype(dtype)
assert byteAlignment >= d.itemsize and byteAlignment % d.itemsize == 0
padding = (byteAlignment - ((ndim0 * d.itemsize) % byteAlignment)) % byteAlignment
N = ndim1 * padding + np.prod(shape) * d.itemsize
tmp = aligned_empty(N, byteAlignment=byteAlignment, dtype=np.uint8, byteOffset=byteOffset).view(dtype=dtype)
bshape = [i for i in shape]
bshape[dim0] = ndim0 + padding // d.itemsize
tmp = tmp.reshape(bshape, order=order)
if tmp.flags['C_CONTIGUOUS']:
tmp = tmp[..., :shape[-1]]
else:
tmp = tmp[:shape[0], ...]
return tmp
def aligned_zeros(shape, byteAlignment=16, dtype=float, byteOffset=0, order='C', alignInnerCoordinate=True):
arr = aligned_empty(shape, dtype=dtype, byteOffset=byteOffset,
order=order, byteAlignment=byteAlignment, alignInnerCoordinate=alignInnerCoordinate)
x = np.zeros((), arr.dtype)
arr[...] = x
return arr
def aligned_ones(shape, byteAlignment=16, dtype=float, byteOffset=0, order='C', alignInnerCoordinate=True):
arr = aligned_empty(shape, dtype=dtype, byteOffset=byteOffset,
order=order, byteAlignment=byteAlignment, alignInnerCoordinate=alignInnerCoordinate)
x = np.ones((), arr.dtype)
arr[...] = x
return arr
......@@ -71,7 +71,7 @@ class SerialDataHandling(DataHandling):
return self._fieldInformation[name]['fSize']
def addArray(self, name, fSize=1, dtype=np.float64, latexName=None, ghostLayers=None, layout=None,
cpu=True, gpu=None):
cpu=True, gpu=None, alignment=False):
if ghostLayers is None:
ghostLayers = self.defaultGhostLayers
if layout is None:
......@@ -99,7 +99,8 @@ class SerialDataHandling(DataHandling):
layoutTuple = spatialLayoutStringToTuple(layout, self.dim)
# cpuArr is always created - since there is no createPycudaArrayWithLayout()
cpuArr = createNumpyArrayWithLayout(layout=layoutTuple, **kwargs)
byteOffset = ghostLayers * np.dtype(dtype).itemsize
cpuArr = createNumpyArrayWithLayout(layout=layoutTuple, alignment=alignment, byteOffset=byteOffset, **kwargs)
cpuArr.fill(np.inf)
if cpu:
......
......@@ -4,6 +4,8 @@ import numpy as np
import sympy as sp
from sympy.core.cache import cacheit
from sympy.tensor import IndexedBase
from pystencils.alignedarray import aligned_empty
from pystencils.data_types import TypedSymbol, createType, createCompositeTypeFromString, StructType
from pystencils.sympyextensions import isIntegerSequence
......@@ -467,11 +469,14 @@ def getLayoutOfArray(arr, indexDimensionIds=[]):
return getLayoutFromStrides(arr.strides, indexDimensionIds)
def createNumpyArrayWithLayout(shape, layout, **kwargs):
def createNumpyArrayWithLayout(shape, layout, alignment=False, byteOffset=0, **kwargs):
"""
Creates a numpy array with
:param shape: shape of the resulting array
:param layout: layout as tuple, where the coordinates are ordered from slow to fast
:param alignment: number of bytes to align the beginning and the innermost coordinate to, or False for no alignment
:param byteOffset: only used when alignment is specified, align not beginning but address at this offset
mostly used to align first inner cell, not ghost cells
>>> res = createNumpyArrayWithLayout(shape=(2, 3, 4, 5), layout=(3, 2, 0, 1))
>>> res.shape
(2, 3, 4, 5)
......@@ -492,7 +497,13 @@ def createNumpyArrayWithLayout(shape, layout, **kwargs):
for a, b in swaps:
shape[a], shape[b] = shape[b], shape[a]
res = np.empty(shape, order='c', **kwargs)
if not alignment:
res = np.empty(shape, order='c', **kwargs)
else:
if alignment is True:
alignment = 8 * 4
res = aligned_empty(shape, alignment, byteOffset=byteOffset)
for a, b in reversed(swaps):
res = res.swapaxes(a, b)
return res
......
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