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

new lbm module: creation functions for kernels

parent 54c2055c
"""
Factory functions for standard LBM methods
"""
import sympy as sp
from copy import copy
from lbmpy.stencils import getStencil
from lbmpy.methods.momentbased import createSRT, createTRT, createOrthogonalMRT
import lbmpy.forcemodels as forceModels
from lbmpy.simplificationfactory import createSimplificationStrategy
from lbmpy.updatekernels import createStreamPullKernel, createPdfArray
from pystencils.cpu.kernelcreation import createKernel as createCpuKernel, addOpenMP
from pystencils.gpucuda.kernelcreation import createCUDAKernel as createGpuKernel
from pystencils.cpu import makePythonFunction as makePythonCpuFunction
from pystencils.gpucuda import makePythonFunction as makePythonGpuFunction
def _getParams(params, optParams):
defaultMethodDescription = {
'target': 'cpu',
'stencil': 'D2Q9',
'method': 'srt', # can be srt, trt or mrt
'relaxationRates': sp.symbols("omega_:10"),
'compressible': False,
'order': 2,
'forceModel': 'none', # can be 'simple', 'luo' or 'guo'
'force': (0, 0, 0),
}
defaultOptimizationDescription = {
'doCseInOpposingDirections': False,
'doOverallCse': False,
'split': True,
'fieldSize': None,
'fieldLayout': 'reverseNumpy', # can be 'numpy' (='c'), 'reverseNumpy' (='f'), 'fzyx', 'zyxf'
'openMP': True,
}
unknownParams = [k for k in params.keys() if k not in defaultMethodDescription]
unknownOptParams = [k for k in optParams.keys() if k not in defaultOptimizationDescription]
if unknownParams:
raise ValueError("Unknown parameter(s): " + ",".join(unknownParams))
if unknownOptParams:
raise ValueError("Unknown optimization parameter(s): " + ",".join(unknownOptParams))
paramsResult = copy(defaultMethodDescription)
paramsResult.update(params)
optParamsResult = copy(defaultOptimizationDescription)
optParamsResult.update(optParams)
return paramsResult, optParamsResult
def createLatticeBoltzmannFunction(optimizationParams={}, **kwargs):
params, optParams = _getParams(kwargs, optimizationParams)
ast = createLatticeBoltzmannKernel(**params, optimizationParams=optParams)
if params['target'] == 'cpu':
if 'openMP' in optParams:
if isinstance(optParams['openMP'], bool) and optParams['openMP']:
addOpenMP(ast)
elif isinstance(optParams['openMP'], int):
addOpenMP(ast, numThreads=optParams['openMP'])
return makePythonCpuFunction(ast)
elif params['target'] == 'gpu':
return makePythonGpuFunction(ast)
else:
return ValueError("'target' has to be either 'cpu' or 'gpu'")
def createLatticeBoltzmannKernel(optimizationParams={}, **kwargs):
params, optParams = _getParams(kwargs, optimizationParams)
updateRule = createLatticeBoltzmannUpdateRule(**params, optimizationParams=optimizationParams)
if params['target'] == 'cpu':
if 'splitGroups' in updateRule.simplificationHints:
splitGroups = updateRule.simplificationHints['splitGroups']
else:
splitGroups = ()
return createCpuKernel(updateRule.allEquations, splitGroups=splitGroups)
elif params['target'] == 'gpu':
return createGpuKernel(updateRule.allEquations)
else:
return ValueError("'target' has to be either 'cpu' or 'gpu'")
def createLatticeBoltzmannUpdateRule(optimizationParams={}, **kwargs):
params, optParams = _getParams(kwargs, optimizationParams)
stencil = getStencil(params['stencil'])
method = createLatticeBoltzmannCollisionRule(**params)
splitInnerLoop = 'split' in optParams and optParams['split']
dirCSE = 'doCseInOpposingDirections'
doCseInOpposingDirections = False if dirCSE not in optParams else optParams[dirCSE]
doOverallCse = False if 'doOverallCse' not in optParams else optParams['doOverallCse']
simplification = createSimplificationStrategy(method, doCseInOpposingDirections, doOverallCse, splitInnerLoop)
collisionRule = simplification(method.getCollisionRule())
if 'fieldSize' in optParams and optParams['fieldSize']:
npField = createPdfArray(optParams['fieldSize'], len(stencil), layout=optParams['fieldLayout'])
updateRule = createStreamPullKernel(collisionRule, numpyField=npField)
else:
layoutName = optParams['fieldLayout']
if layoutName == 'fzyx' or 'zyxf':
dim = len(stencil[0])
layoutName = tuple(reversed(range(dim)))
updateRule = createStreamPullKernel(collisionRule, genericLayout=layoutName)
return updateRule
def createLatticeBoltzmannCollisionRule(**params):
params, _ = _getParams(params, {})
stencil = getStencil(params['stencil'])
dim = len(stencil[0])
if 'forceModel' in params:
forceModelName = params['forceModel']
if forceModelName.lower() == 'none':
forceModel = None
elif forceModelName.lower() == 'simple':
forceModel = forceModels.Simple(params['force'][:dim])
elif forceModelName.lower() == 'luo':
forceModel = forceModels.Luo(params['force'][:dim])
elif forceModelName.lower() == 'guo':
forceModel = forceModels.Guo(params['force'][:dim])
else:
raise ValueError("Unknown force model %s" % (forceModelName,))
else:
forceModel = None
commonParams = {
'compressible': params['compressible'],
'equilibriumAccuracyOrder': params['order'],
'forceModel': forceModel,
}
methodName = params['method']
relaxationRates = params['relaxationRates']
if methodName.lower() == 'srt':
assert len(relaxationRates) >= 1, "Not enough relaxation rates"
method = createSRT(stencil, relaxationRates[0], **commonParams)
elif methodName.lower() == 'trt':
assert len(relaxationRates) >= 2, "Not enough relaxation rates"
method = createTRT(stencil, relaxationRates[0], relaxationRates[1], **commonParams)
elif methodName.lower() == 'mrt':
nextRelaxationRate = 0
def relaxationRateGetter(momentGroup):
nonlocal nextRelaxationRate
res = relaxationRates[nextRelaxationRate]
nextRelaxationRate += 1
return res
method = createOrthogonalMRT(stencil, relaxationRateGetter, **commonParams)
else:
raise ValueError("Unknown method %s" % (methodName,))
return method
import sympy as sp
from collections import defaultdict
from pystencils import Field
def createLbmSplitGroups(lbmCollisionEqs):
......@@ -18,16 +19,27 @@ def createLbmSplitGroups(lbmCollisionEqs):
sh = lbmCollisionEqs.simplificationHints
assert 'velocity' in sh, "Needs simplification hint 'velocity': Sequence of velocity symbols"
pdfSymbols = lbmCollisionEqs.method.preCollisionPdfSymbols
preCollisionSymbols = set(lbmCollisionEqs.method.preCollisionPdfSymbols)
nonCenterPostCollisionSymbols = set(lbmCollisionEqs.method.postCollisionPdfSymbols[1:])
postCollisionSymbols = set(lbmCollisionEqs.method.postCollisionPdfSymbols)
stencil = lbmCollisionEqs.method.stencil
importantSubExpressions = {e.lhs for e in lbmCollisionEqs.subexpressions
if pdfSymbols.intersection(lbmCollisionEqs.getDependentSymbols([e.lhs]))}
for eq in lbmCollisionEqs.mainEquations[1:]:
if preCollisionSymbols.intersection(lbmCollisionEqs.getDependentSymbols([e.lhs]))}
otherWrittenFields = []
for eq in lbmCollisionEqs.mainEquations:
if eq.lhs not in postCollisionSymbols and isinstance(eq.lhs, Field.Access):
otherWrittenFields.append(eq.lhs)
if eq.lhs not in nonCenterPostCollisionSymbols:
continue
importantSubExpressions.intersection_update(eq.rhs.atoms(sp.Symbol))
subexpressionsToPreCompute = list(sh['velocity']) + list(importantSubExpressions)
splitGroups = [subexpressionsToPreCompute, ]
importantSubExpressions.update(sh['velocity'])
subexpressionsToPreCompute = list(importantSubExpressions)
splitGroups = [subexpressionsToPreCompute + otherWrittenFields, ]
directionGroups = defaultdict(list)
dim = len(stencil[0])
......@@ -45,4 +57,5 @@ def createLbmSplitGroups(lbmCollisionEqs):
directionGroups[direction].append(eq.lhs)
splitGroups += directionGroups.values()
return splitGroups
lbmCollisionEqs.simplificationHints['splitGroups'] = splitGroups
return lbmCollisionEqs
......@@ -380,20 +380,20 @@ def createOrthogonalMRT(stencil, relaxationRateGetter=None, compressible=False,
sq = x ** 2 + y ** 2 + z ** 2
nestedMoments = [
[one, x, y, z], # [0, 3, 5, 7]
[3 * x ** 2 - sq, y ** 2 - z ** 2, x * y, y * z, x * z], # [9, 10, 11, 12, 13]
[sq - 1], # [1]
[3 * sq ** 2 - 6 * sq + 1], # [2]
[(3 * sq - 5) * x, (3 * sq - 5) * y, (3 * sq - 5) * z], # [4, 6, 8]
[3 * x ** 2 - sq, y ** 2 - z ** 2, x * y, y * z, x * z], # [9, 10, 11, 12, 13]
[x * y * z]
]
elif stencilsHaveSameEntries(stencil, getStencil("D3Q19")):
sq = x ** 2 + y ** 2 + z ** 2
nestedMoments = [
[one, x, y, z], # [0, 3, 5, 7]
[3 * x ** 2 - sq, y ** 2 - z ** 2, x * y, y * z, x * z], # [9, 11, 13, 14, 15]
[sq - 1], # [1]
[3 * sq ** 2 - 6 * sq + 1], # [2]
[(3 * sq - 5) * x, (3 * sq - 5) * y, (3 * sq - 5) * z], # [4, 6, 8]
[3 * x ** 2 - sq, y ** 2 - z ** 2, x * y, y * z, x * z], # [9, 11, 13, 14, 15]
[(2 * sq - 3) * (3 * x ** 2 - sq), (2 * sq - 3) * (y ** 2 - z ** 2)], # [10, 12]
[(y ** 2 - z ** 2) * x, (z ** 2 - x ** 2) * y, (x ** 2 - y ** 2) * z] # [16, 17, 18]
]
......
from functools import partial
import sympy as sp
from lbmpy.innerloopsplit import createLbmSplitGroups
from pystencils.equationcollection.simplifications import applyOnAllEquations, \
subexpressionSubstitutionInMainEquations, sympyCSE, addSubexpressionsForDivisions
def createSimplificationStrategy(lbmMethod, doCseInOpposingDirections=False, doOverallCse=False):
def createSimplificationStrategy(lbmMethod, doCseInOpposingDirections=False, doOverallCse=False, splitInnerLoop=False):
from pystencils.equationcollection import SimplificationStrategy
from lbmpy.methods import MomentBasedLbmMethod
from lbmpy.methods.momentbasedsimplifications import replaceSecondOrderVelocityProducts, \
......@@ -25,6 +27,8 @@ def createSimplificationStrategy(lbmMethod, doCseInOpposingDirections=False, doO
s.add(replaceCommonQuadraticAndConstantTerm)
s.add(factorDensityAfterFactoringRelaxationTimes)
s.add(subexpressionSubstitutionInMainEquations)
if splitInnerLoop:
s.add(createLbmSplitGroups)
if doCseInOpposingDirections:
s.add(cseInOpposingDirections)
......@@ -34,16 +38,3 @@ def createSimplificationStrategy(lbmMethod, doCseInOpposingDirections=False, doO
s.add(addSubexpressionsForDivisions)
return s
if __name__ == '__main__':
from lbmpy.stencils import getStencil
from lbmpy.methods.momentbased import createOrthogonalMRT
stencil = getStencil("D2Q9")
m = createOrthogonalMRT(stencil, compressible=True)
cr = m.getCollisionRule()
simp = createSimplificationStrategy(m)
simp(cr)
import numpy as np
from pystencils import Field
from pystencils.sympyextensions import fastSubs
from lbmpy.fieldaccess import StreamPullTwoFieldsAccessor
......@@ -18,7 +19,7 @@ def createLBMKernel(collisionRule, inputField, outputField, accessor):
:return: LbmCollisionRule where pre- and post collision symbols have been replaced
"""
method = collisionRule.method
preCollisionSymbols = method.preCollisionSymbols
preCollisionSymbols = method.preCollisionPdfSymbols
postCollisionSymbols = method.postCollisionPdfSymbols
substitutions = {}
......@@ -29,7 +30,15 @@ def createLBMKernel(collisionRule, inputField, outputField, accessor):
substitutions[preCollisionSymbols[idx]] = inputAccess
substitutions[postCollisionSymbols[idx]] = outputAccess
return collisionRule.copyWithSubstitutionsApplied(substitutions)
result = collisionRule.copyWithSubstitutionsApplied(substitutions)
if 'splitGroups' in result.simplificationHints:
newSplitGroups = []
for splitGroup in result.simplificationHints['splitGroups']:
newSplitGroups.append([fastSubs(e, substitutions) for e in splitGroup])
result.simplificationHints['splitGroups'] = newSplitGroups
return result
def createStreamPullKernel(collisionRule, numpyField=None, srcFieldName="src", dstFieldName="dst",
......@@ -60,6 +69,36 @@ def createStreamPullKernel(collisionRule, numpyField=None, srcFieldName="src", d
return createLBMKernel(collisionRule, src, dst, StreamPullTwoFieldsAccessor)
# ---------------------------------- Pdf array creation for various layouts --------------------------------------------
def createPdfArray(size, numDirections, ghostLayers=1, layout='fzyx'):
"""
Creates an empy numpy array for a pdf field with the specified memory layout.
Examples:
>>> createPdfArray((3, 4, 5), 9, layout='zyxf', ghostLayers=0).shape
(3, 4, 5, 9)
>>> createPdfArray((3, 4, 5), 9, layout='zyxf', ghostLayers=0).strides
(72, 216, 864, 8)
>>> createPdfArray((3, 4), 9, layout='zyxf', ghostLayers=1).shape
(5, 6, 9)
>>> createPdfArray((3, 4), 9, layout='zyxf', ghostLayers=1).strides
(72, 360, 8)
"""
sizeWithGl = [s + 2 * ghostLayers for s in size]
if layout == "fzyx" or layout == 'f' or layout == 'reverseNumpy':
return np.empty(sizeWithGl + [numDirections], order='f')
elif layout == 'c' or layout == 'numpy':
return np.empty(sizeWithGl + [numDirections], order='c')
elif layout == 'zyxf':
res = np.empty(list(reversed(sizeWithGl)) + [numDirections], order='c')
res = res.swapaxes(0, 1)
if len(size) == 3:
res = res.swapaxes(1, 2)
res = res.swapaxes(0, 1)
return res
# ------------------------------------------- Add output fields to kernel ----------------------------------------------
......
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