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

Restructuring: moved to pystencils

parent 3414de0a
from collections import defaultdict
import sympy as sp
from pystencils.generator import resolveFieldAccesses
from pystencils.generator import typeAllEquations, Block, KernelFunction, parseBasePointerInfo
BLOCK_IDX = list(sp.symbols("blockIdx.x blockIdx.y blockIdx.z"))
THREAD_IDX = list(sp.symbols("threadIdx.x threadIdx.y threadIdx.z"))
GPU Access Patterns
- knows about the iteration range
- know about mapping of field indices to CUDA block and thread indices
- iterates over spatial coordinates - constructed with a specific number of coordinates
- can
def getLinewiseCoordinateAccessExpression(field, indexCoordinate):
availableIndices = [THREAD_IDX[0]] + BLOCK_IDX
fastestCoordinate = field.layout[-1]
availableIndices[fastestCoordinate], availableIndices[0] = availableIndices[0], availableIndices[fastestCoordinate]
cudaIndices = availableIndices[:field.spatialDimensions]
offsetToCell = sum([cudaIdx * stride for cudaIdx, stride in zip(cudaIndices, field.spatialStrides)])
indexOffset = sum([idx * indexStride for idx, indexStride in zip(indexCoordinate, field.indexStrides)])
return sp.simplify(offsetToCell + indexOffset)
def getLinewiseCoordinates(field):
availableIndices = [THREAD_IDX[0]] + BLOCK_IDX
d = field.spatialDimensions + field.indexDimensions
fastestCoordinate = field.layout[-1]
result = availableIndices[:d]
result[0], result[fastestCoordinate] = result[fastestCoordinate], result[0]
return result
def createCUDAKernel(listOfEquations, functionName="kernel", typeForSymbol=defaultdict(lambda: "double")):
fieldsRead, fieldsWritten, assignments = typeAllEquations(listOfEquations, typeForSymbol)
for f in fieldsRead - fieldsWritten:
code = KernelFunction(Block(assignments), functionName)
code.qualifierPrefix = "__global__ "
code.variablesToIgnore.update(BLOCK_IDX + THREAD_IDX)
coordMapping = getLinewiseCoordinates(list(fieldsRead)[0])
allFields = fieldsRead.union(fieldsWritten)
basePointerInfo = [['spatialInner0']]
basePointerInfos = { parseBasePointerInfo(basePointerInfo, [0, 1, 2], f) for f in allFields}
resolveFieldAccesses(code, fieldToFixedCoordinates={'src': coordMapping, 'dst': coordMapping},
return code
if __name__ == "__main__":
import sympy as sp
from lbmpy.stencils import getStencil
from lbmpy.collisionoperator import makeSRT
from lbmpy.lbmgenerator import createLbmEquations
latticeModel = makeSRT(getStencil("D2Q9"), order=2, compressible=False)
r = createLbmEquations(latticeModel, doCSE=True)
kernel = createCUDAKernel(r)
from pycuda.compiler import SourceModule
mod = SourceModule(str(kernel.generateC()))
func = mod.get_function("kernel")
import numpy as np
import pycuda.driver as cuda
from pycuda.compiler import SourceModule
def numpyTypeFromString(typename, includePointers=True):
import ctypes as ct
typename = typename.replace("*", " * ")
typeComponents = typename.split()
basicTypeMap = {
'double': np.float64,
'float': np.float32,
'int': np.int32,
'long': np.int64,
resultType = None
for typeComponent in typeComponents:
typeComponent = typeComponent.strip()
if typeComponent == "const" or typeComponent == "restrict" or typeComponent == "volatile":
if typeComponent in basicTypeMap:
resultType = basicTypeMap[typeComponent]
elif typeComponent == "*" and includePointers:
assert resultType is not None
resultType = ct.POINTER(resultType)
return resultType
def buildNumpyArgumentList(kernelFunctionNode, argumentDict):
result = []
for arg in kernelFunctionNode.parameters:
if arg.isFieldArgument:
field = argumentDict[arg.fieldName]
if arg.isFieldPtrArgument:
elif arg.isFieldShapeArgument:
strideArr = np.array(field.strides, dtype=np.int32) / field.dtype.itemsize
elif arg.isFieldStrideArgument:
shapeArr = np.array(field.shape, dtype=np.int32)
assert False
param = argumentDict[]
expectedType = numpyTypeFromString(arg.dtype)
return result
def makePythonFunction(kernelFunctionNode, argumentDict):
mod = SourceModule(str(kernelFunctionNode.generateC()))
func = mod.get_function(kernelFunctionNode.functionName)
# 1) get argument list
args = buildNumpyArgumentList(kernelFunctionNode, argumentDict)
# 2) determine block and grid tuples
# TODO prepare the function here
from itertools import chain
import numpy as np
import sympy as sp
from sympy.core.cache import cacheit
from sympy.tensor import IndexedBase
from pystencils.typedsymbol import TypedSymbol
def getLayoutFromNumpyArray(arr):
Returns a list indicating the memory layout (linearization order) of the numpy array.
>>> getLayoutFromNumpyArray(np.zeros([3,3,3]))
[0, 1, 2]
In this example the loop over the zeroth coordinate should be the outermost loop,
followed by the first and second. Elements arr[x,y,0] and arr[x,y,1] are adjacent in memory.
Normally constructed numpy arrays have this order, however by stride tricks or other frameworks, arrays
with different memory layout can be created.
coordinates = list(range(len(arr.shape)))
return [x for (y, x) in sorted(zip(arr.strides, coordinates), key=lambda pair: pair[0], reverse=True)]
def numpyDataTypeToC(dtype):
"""Mapping numpy data types to C data types"""
if dtype == np.float64:
return "double"
elif dtype == np.float32:
return "float"
elif dtype == np.int32:
return "int"
raise NotImplementedError()
def offsetComponentToDirectionString(coordinateId, value):
Translates numerical offset to string notation.
x offsets are labeled with east 'E' and 'W',
y offsets with north 'N' and 'S' and
z offsets with top 'T' and bottom 'B'
If the absolute value of the offset is bigger than 1, this number is prefixed.
:param coordinateId: integer 0, 1 or 2 standing for x,y and z
:param value: integer offset
>>> offsetComponentToDirectionString(0, 1)
>>> offsetComponentToDirectionString(1, 2)
nameComponents = (('W', 'E'), # west, east
('S', 'N'), # south, north
('B', 'T'), # bottom, top
if value == 0:
result = ""
elif value < 0:
result = nameComponents[coordinateId][0]
result = nameComponents[coordinateId][1]
if abs(value) > 1:
result = "%d%s" % (abs(value), result)
return result
def offsetToDirectionString(offsetTuple):
Translates numerical offset to string notation.
For details see :func:`offsetComponentToDirectionString`
:param offsetTuple: 3-tuple with x,y,z offset
>>> offsetToDirectionString([1, -1, 0])
>>> offsetToDirectionString(([-3, 0, -2]))
names = ["", "", ""]
for i in range(len(offsetTuple)):
names[i] = offsetComponentToDirectionString(i, offsetTuple[i])
name = "".join(reversed(names))
if name == "":
name = "C"
return name
def directionStringToOffset(directionStr, dim=3):
Reverse mapping of :func:`offsetToDirectionString`
:param directionStr: string representation of offset
:param dim: dimension of offset, i.e the length of the returned list
>>> directionStringToOffset('NW', dim=3)
array([-1, 1, 0])
>>> directionStringToOffset('NW', dim=2)
array([-1, 1])
>>> directionStringToOffset(offsetToDirectionString([3,-2,1]))
array([ 3, -2, 1])
offsetMap = {
'C': np.array([0, 0, 0]),
'W': np.array([-1, 0, 0]),
'E': np.array([1, 0, 0]),
'S': np.array([0, -1, 0]),
'N': np.array([0, 1, 0]),
'B': np.array([0, 0, -1]),
'T': np.array([0, 0, 1]),
offset = np.array([0, 0, 0])
while len(directionStr) > 0:
factor = 1
firstNonDigit = 0
while directionStr[firstNonDigit].isdigit():
firstNonDigit += 1
if firstNonDigit > 0:
factor = int(directionStr[:firstNonDigit])
directionStr = directionStr[firstNonDigit:]
curOffset = offsetMap[directionStr[0]]
offset += factor * curOffset
directionStr = directionStr[1:]
return offset[:dim]
class Field:
With fields one can formulate stencil-like update rules on structured grids.
This Field class knows about the dimension, memory layout (strides) and optionally about the size of an array.
To create a field use one of the static create* members. There are two options:
1. create a kernel with fixed loop sizes i.e. the shape of the array is already known. This is usually the
case if just-in-time compilation directly from Python is done. (see Field.createFromNumpyArray)
2. create a more general kernel that works for variable array sizes. This can be used to create kernels
beforehand for a library. (see Field.createGeneric)
A field has spatial and index dimensions, where the spatial dimensions come first.
The interpretation is that the field has multiple cells in (usually) two or three dimensional space which are
looped over. Additionally N values are stored per cell. In this case spatialDimensions is two or three,
and indexDimensions equals N. If you want to store a matrix on each point in a two dimensional grid, there
are four dimensions, two spatial and two index dimensions. len(arr.shape) == spatialDims + indexDims
When accessing (indexing) a field the result is a FieldAccess which is derived from sympy Symbol.
First specify the spatial offsets in [], then in case indexDimension>0 the indices in ()
e.g. f[-1,0,0](7)
Example without index dimensions:
>>> a = np.zeros([10, 10])
>>> f = Field.createFromNumpyArray("f", a, indexDimensions=0)
>>> jacobi = ( f[-1,0] + f[1,0] + f[0,-1] + f[0,1] ) / 4
Example with index dimensions: LBM D2Q9 stream pull
>>> stencil = np.array([[0,0], [0,1], [0,-1]])
>>> src = Field.createGeneric("src", spatialDimensions=2, indexDimensions=1)
>>> dst = Field.createGeneric("dst", spatialDimensions=2, indexDimensions=1)
>>> for i, offset in enumerate(stencil):
... sp.Eq(dst[0,0](i), src[-offset](i))
Eq(dst_C^0, src_C^0)
Eq(dst_C^1, src_S^1)
Eq(dst_C^2, src_N^2)
def createFromNumpyArray(fieldName, npArray, indexDimensions=0):
Creates a field based on the layout, data type, and shape of a given numpy array.
Kernels created for these kind of fields can only be called with arrays of the same layout, shape and type.
:param fieldName: symbolic name for the field
:param npArray: numpy array
:param indexDimensions: see documentation of Field
spatialDimensions = len(npArray.shape) - indexDimensions
if spatialDimensions < 1:
raise ValueError("Too many index dimensions. At least one spatial dimension required")
fullLayout = getLayoutFromNumpyArray(npArray)
spatialLayout = tuple([i for i in fullLayout if i < spatialDimensions])
assert len(spatialLayout) == spatialDimensions
strides = tuple([s // np.dtype(npArray.dtype).itemsize for s in npArray.strides])
shape = tuple([int(s) for s in npArray.shape])
return Field(fieldName, npArray.dtype, spatialLayout, shape, strides)
def createGeneric(fieldName, spatialDimensions, dtype=np.float64, indexDimensions=0, layout=None):
Creates a generic field where the field size is not fixed i.e. can be called with arrays of different sizes
:param fieldName: symbolic name for the field
:param dtype: numpy data type of the array the kernel is called with later
:param spatialDimensions: see documentation of Field
:param indexDimensions: see documentation of Field
:param layout: tuple specifying the loop ordering of the spatial dimensions e.g. (2, 1, 0 ) means that
the outer loop loops over dimension 2, the second outer over dimension 1, and the inner loop
over dimension 0
if not layout:
layout = tuple(reversed(range(spatialDimensions)))
if len(layout) != spatialDimensions:
raise ValueError("Layout")
shapeSymbol = IndexedBase(TypedSymbol(Field.SHAPE_PREFIX + fieldName, Field.SHAPE_DTYPE), shape=(1,))
strideSymbol = IndexedBase(TypedSymbol(Field.STRIDE_PREFIX + fieldName, Field.STRIDE_DTYPE), shape=(1,))
totalDimensions = spatialDimensions + indexDimensions
shape = tuple([shapeSymbol[i] for i in range(totalDimensions)])
strides = tuple([strideSymbol[i] for i in range(totalDimensions)])
return Field(fieldName, dtype, layout, shape, strides)
def __init__(self, fieldName, dtype, layout, shape, strides):
"""Do not use directly. Use static create* methods"""
self._fieldName = fieldName
self._dtype = numpyDataTypeToC(dtype)
self._layout = layout
self._shape = shape
self._strides = strides
self._readonly = False
def spatialDimensions(self):
return len(self._layout)
def indexDimensions(self):
return len(self._shape) - len(self._layout)
def layout(self):
return self._layout
def name(self):
return self._fieldName
def shape(self):
return self._shape
def spatialShape(self):
return self._shape[:self.spatialDimensions]
def indexShape(self):
return self._shape[self.spatialDimensions:]
def spatialStrides(self):
return self._strides[:self.spatialDimensions]
def indexStrides(self):
return self._strides[self.spatialDimensions:]
def strides(self):
return self._strides
def dtype(self):
return self._dtype
def readOnly(self):
return self._readonly
def setReadOnly(self, value=True):
self._readonly = value
def __repr__(self):
return self._fieldName
def __getitem__(self, offset):
if type(offset) is np.ndarray:
offset = tuple(offset)
if type(offset) is str:
offset = tuple(directionStringToOffset(offset, self.spatialDimensions))
if type(offset) is not tuple:
offset = (offset,)
if len(offset) != self.spatialDimensions:
raise ValueError("Wrong number of spatial indices: "
"Got %d, expected %d" % (len(offset), self.spatialDimensions))
return Field.Access(self, offset)
def __call__(self, *args, **kwargs):
center = tuple([0]*self.spatialDimensions)
return Field.Access(self, center)(*args, **kwargs)
def __hash__(self):
return hash((self._layout, self._shape, self._strides, self._dtype, self._fieldName))
def __eq__(self, other):
selfTuple = (self.shape, self.strides,, self.dtype)
otherTuple = (other.shape, other.strides,, other.dtype)
return selfTuple == otherTuple
PREFIX = "f"
STRIDE_DTYPE = "const int *"
SHAPE_DTYPE = "const int *"
class Access(sp.Symbol):
def __new__(cls, name, *args, **kwargs):
obj = Field.Access.__xnew_cached_(cls, name, *args, **kwargs)
return obj
def __new_stage2__(self, field, offsets=(0, 0, 0), idx=None):
fieldName =
offsetsAndIndex = chain(offsets, idx) if idx is not None else offsets
constantOffsets = not any([isinstance(o, sp.Basic) for o in offsetsAndIndex])
if not idx:
idx = tuple([0] * field.indexDimensions)
if constantOffsets:
offsetName = offsetToDirectionString(offsets)
if field.indexDimensions == 0:
obj = super(Field.Access, self).__xnew__(self, fieldName + "_" + offsetName)
elif field.indexDimensions == 1:
obj = super(Field.Access, self).__xnew__(self, fieldName + "_" + offsetName + "^" + str(idx[0]))
idxStr = ",".join([str(e) for e in idx])
obj = super(Field.Access, self).__xnew__(self, fieldName + "_" + offsetName + "^" + idxStr)
offsetName = "%0.10X" % (abs(hash(tuple(offsetsAndIndex))))
obj = super(Field.Access, self).__xnew__(self, fieldName + "_" + offsetName)
obj._field = field
obj._offsets = []
for o in offsets:
if isinstance(o, sp.Basic):
obj._offsetName = offsetName
obj._index = idx
return obj
__xnew__ = staticmethod(__new_stage2__)
__xnew_cached_ = staticmethod(cacheit(__new_stage2__))
def __call__(self, *idx):
if self._index != tuple([0]*self.field.indexDimensions):
print(self._index, tuple([0]*self.field.indexDimensions))
raise ValueError("Indexing an already indexed Field.Access")
idx = tuple(idx)
if len(idx) != self.field.indexDimensions and idx != (0,):
raise ValueError("Wrong number of indices: "
"Got %d, expected %d" % (len(idx), self.field.indexDimensions))
return Field.Access(self.field, self._offsets, idx)
def field(self):
return self._field
def offsets(self):
return self._offsets
def requiredGhostLayers(self):
return int(np.max(np.abs(self._offsets)))
def nrOfCoordinates(self):
return len(self._offsets)
def offsetName(self):
return self._offsetName
def index(self):
return self._index
def _hashable_content(self):
superClassContents = list(super(Field.Access, self)._hashable_content())
t = tuple([*superClassContents, hash(self._field), self._index] + self._offsets)
return t
import sympy as sp
import numpy as np
from lbmpy.generator import Field
import sympy as sp
from pystencils.generator import Field
def __upDownOffsets(d, dim):
This diff is collapsed.
import os
import subprocess
from ctypes import cdll, c_double, c_float, sizeof
from tempfile import TemporaryDirectory
import numpy as np
'compiler': 'g++',
'flags': '-Ofast -DNDEBUG -fPIC -shared -march=native -fopenmp',
'compiler': '/software/intel/2017/bin/icpc',
'flags': '-Ofast -DNDEBUG -fPIC -shared -march=native -fopenmp -Wl,-rpath=/software/intel/2017/lib/intel64',
'env': {
'LM_PROJECT': 'iwia',
'compiler': 'clang++',
'flags': '-Ofast -DNDEBUG -fPIC -shared -march=native -fopenmp',
def ctypeFromString(typename, includePointers=True):
import ctypes as ct
typename = typename.replace("*", " * ")
typeComponents = typename.split()
basicTypeMap = {
'double': ct.c_double,
'float': ct.c_float,
'int': ct.c_int,
'long': ct.c_long,
resultType = None
for typeComponent in typeComponents:
typeComponent = typeComponent.strip()
if typeComponent == "const" or typeComponent == "restrict" or typeComponent == "volatile":
if typeComponent in basicTypeMap:
resultType = basicTypeMap[typeComponent]
elif typeComponent == "*" and includePointers:
assert resultType is not None
resultType = ct.POINTER(resultType)
return resultType
def ctypeFromNumpyType(numpyType):
typeMap = {
np.dtype('float64'): c_double,
np.dtype('float32'): c_float,
return typeMap[numpyType]
def compileAndLoad(kernelFunctionNode):
with TemporaryDirectory() as tmpDir:
srcFile = os.path.join(tmpDir, 'source.cpp')
with open(srcFile, 'w') as sourceFile:
print('#include <iostream>', file=sourceFile)
print("#include <cmath>", file=sourceFile)
print('extern "C" { ', file=sourceFile)
print(kernelFunctionNode.generateC(), file=sourceFile)
print('}', file=sourceFile)
compilerCmd = [CONFIG['compiler']] + CONFIG['flags'].split()
libFile = os.path.join(tmpDir, "")
compilerCmd += [srcFile, '-o', libFile]
configEnv = CONFIG['env'] if 'env' in CONFIG else {}
env = os.environ.copy()
env.update(configEnv), env=env)
showAssembly = False
if showAssembly:
assemblyFile = os.path.join(