Commit 5bed1ae5 authored by Martin Bauer's avatar Martin Bauer
Browse files

New code generation interface for waLBerla

parent 5e6fe078
from .walberla_lbm_generation import generate_lattice_model_files, generate_lattice_model, RefinementScaling, \
create_lb_method
from pystencils import Field
from .walberla_lbm_generation import generate_lattice_model, RefinementScaling
from .boundary import generate_boundary
__all__ = ['Field', 'generate_lattice_model_files', 'generate_lattice_model', 'RefinementScaling', 'create_lb_method']
__all__ = ['generate_lattice_model', 'RefinementScaling', 'generate_boundary']
......@@ -3,6 +3,7 @@ from jinja2 import Environment, PackageLoader
from pystencils import Field, FieldType
from pystencils.data_types import create_type, TypedSymbol
from pystencils_walberla.codegen import default_create_kernel_parameters
from pystencils_walberla.jinja_filters import add_pystencils_filters_to_jinja_env
from lbmpy.boundaries.boundaryhandling import create_lattice_boltzmann_boundary_kernel
......@@ -11,44 +12,24 @@ from pystencils.boundaries.createindexlist import numpy_data_type_for_boundary_o
from lbmpy_walberla.walberla_lbm_generation import KernelInfo
def struct_from_numpy_dtype(struct_name, numpy_dtype):
result = "struct %s { \n" % (struct_name,)
equality_compare = []
constructor_params = []
constructor_initializer_list = []
for name, (sub_type, offset) in numpy_dtype.fields.items():
pystencils_type = create_type(sub_type)
result += " %s %s;\n" % (pystencils_type, name)
if name in boundary_index_array_coordinate_names or name == direction_member_name:
constructor_params.append("%s %s_" % (pystencils_type, name))
constructor_initializer_list.append("%s(%s_)" % (name, name))
else:
constructor_initializer_list.append("%s()" % name)
if pystencils_type.is_float():
equality_compare.append("floatIsEqual(%s, o.%s)" % (name, name))
else:
equality_compare.append("%s == o.%s" % (name, name))
result += " %s(%s) : %s {}\n" % \
(struct_name, ", ".join(constructor_params), ", ".join(constructor_initializer_list))
result += " bool operator==(const %s & o) const {\n return %s;\n }\n" % \
(struct_name, " && ".join(equality_compare))
result += "};\n"
return result
def generate_boundary(generation_context, class_name, boundary_object, lb_method, **create_kernel_params):
struct_name = "IndexInfo"
boundary_object.name = class_name
create_kernel_params = default_create_kernel_parameters(generation_context, create_kernel_params)
target = create_kernel_params['target']
def create_boundary_class(boundary_object, lb_method, double_precision=True, target='cpu'):
struct_name = "IndexInfo"
index_struct_dtype = numpy_data_type_for_boundary_object(boundary_object, lb_method.dim)
pdf_field = Field.create_generic('pdfs', lb_method.dim, np.float64 if double_precision else np.float32,
pdf_field = Field.create_generic('pdfs', lb_method.dim,
np.float64 if generation_context.double_accuracy else np.float32,
index_dimensions=1, layout='fzyx', index_shape=[len(lb_method.stencil)])
index_field = Field('indexVector', FieldType.INDEXED, index_struct_dtype, layout=[0],
shape=(TypedSymbol("indexVectorSize", create_type(np.int64)), 1), strides=(1, 1))
kernel = create_lattice_boltzmann_boundary_kernel(pdf_field, index_field, lb_method, boundary_object, target=target)
kernel = create_lattice_boltzmann_boundary_kernel(pdf_field, index_field, lb_method, boundary_object, target=target,
openmp=generation_context.openmp)
kernel.function_name = "boundary_" + boundary_object.name
stencil_info = [(i, ", ".join([str(e) for e in d])) for i, d in enumerate(lb_method.stencil)]
......@@ -67,18 +48,36 @@ def create_boundary_class(boundary_object, lb_method, double_precision=True, tar
env = Environment(loader=PackageLoader('lbmpy_walberla'))
add_pystencils_filters_to_jinja_env(env)
header_file = env.get_template('Boundary.tmpl.h').render(**context)
cpp_file = env.get_template('Boundary.tmpl.cpp').render(**context)
return header_file, cpp_file
header = env.get_template('Boundary.tmpl.h').render(**context)
source = env.get_template('Boundary.tmpl.cpp').render(**context)
source_extension = "cpp" if create_kernel_params.get("target", "cpu") == "cpu" else "cu"
generation_context.write_file("{}.h".format(class_name), header)
generation_context.write_file("{}.{}".format(class_name, source_extension), source)
if __name__ == '__main__':
from lbmpy.boundaries import UBB
from lbmpy.creationfunctions import create_lb_method
boundary = UBB(lambda: 0, dim=2)
method = create_lb_method(stencil='D2Q9', method='srt')
header, source = create_boundary_class(boundary, method)
print(source)
def struct_from_numpy_dtype(struct_name, numpy_dtype):
result = "struct %s { \n" % (struct_name,)
equality_compare = []
constructor_params = []
constructor_initializer_list = []
for name, (sub_type, offset) in numpy_dtype.fields.items():
pystencils_type = create_type(sub_type)
result += " %s %s;\n" % (pystencils_type, name)
if name in boundary_index_array_coordinate_names or name == direction_member_name:
constructor_params.append("%s %s_" % (pystencils_type, name))
constructor_initializer_list.append("%s(%s_)" % (name, name))
else:
constructor_initializer_list.append("%s()" % name)
if pystencils_type.is_float():
equality_compare.append("floatIsEqual(%s, o.%s)" % (name, name))
else:
equality_compare.append("%s == o.%s" % (name, name))
result += " %s(%s) : %s {}\n" % \
(struct_name, ", ".join(constructor_params), ", ".join(constructor_initializer_list))
result += " bool operator==(const %s & o) const {\n return %s;\n }\n" % \
(struct_name, " && ".join(equality_compare))
result += "};\n"
return result
import sympy as sp
from sympy.tensor import IndexedBase
from jinja2 import Environment, PackageLoader, Template
import os
import numpy as np
import inspect
from lbmpy.stencils import get_stencil
from pystencils import AssignmentCollection
from pystencils.astnodes import SympyAssignment
from pystencils.sympyextensions import get_symmetric_part
from pystencils.field import Field
from pystencils.stencils import offset_to_direction_string
from pystencils.stencils import offset_to_direction_string, stencils_have_same_entries
from pystencils.backends.cbackend import CustomSympyPrinter, CBackend
from pystencils.data_types import TypedSymbol, cast_func, create_type
from pystencils.data_types import TypedSymbol
from pystencils.transformations import add_types
from pystencils_walberla.sweep import KernelInfo
from pystencils_walberla.codegen import KernelInfo, default_create_kernel_parameters
from pystencils_walberla.jinja_filters import add_pystencils_filters_to_jinja_env
from pystencils.cpu import add_openmp, create_kernel
from pystencils import create_kernel
from lbmpy.relaxationrates import relaxation_rate_scaling
from lbmpy.creationfunctions import create_lb_method, update_with_default_parameters,\
create_lb_ast
from lbmpy.creationfunctions import update_with_default_parameters, create_lb_update_rule
from lbmpy.updatekernels import create_stream_pull_only_kernel
cpp_printer = CustomSympyPrinter()
REFINEMENT_SCALE_FACTOR = sp.Symbol("level_scale_factor")
def stencil_switch_statement(stencil, values):
template = Template("""
using namespace stencil;
switch( direction ) {
{% for direction_name, value in dir_to_value_dict.items() -%}
case {{direction_name}}: return {{value}};
{% endfor -%}
default:
WALBERLA_ABORT("Invalid Direction");
}
""")
dir_to_value_dict = {offset_to_direction_string(d): cpp_printer.doprint(v) for d, v in zip(stencil, values)}
return template.render(dir_to_value_dict=dir_to_value_dict)
def field_and_symbol_substitute(expr, variable_prefix="lm.", variables_without_prefix=[]):
variables_without_prefix = [v.name if isinstance(v, sp.Symbol) else v for v in variables_without_prefix]
substitutions = {}
for sym in expr.atoms(sp.Symbol):
if isinstance(sym, Field.Access):
fa = sym
prefix = "" if fa.field.name in variables_without_prefix else variable_prefix
if fa.field.index_dimensions == 0:
substitutions[fa] = sp.Symbol("%s%s->get(x,y,z)" % (prefix, fa.field.name))
else:
assert fa.field.index_dimensions == 1, "walberla supports only 0 or 1 index dimensions"
substitutions[fa] = sp.Symbol("%s%s->get(x,y,z,%s)" % (prefix, fa.field.name, fa.index[0]))
else:
if sym.name not in variables_without_prefix:
substitutions[sym] = sp.Symbol(variable_prefix + sym.name)
return expr.subs(substitutions)
def expression_to_code(expr, variable_prefix="lm.", variables_without_prefix=[]):
"""
Takes a sympy expression and creates a C code string from it. Replaces field accesses by
walberla field accesses i.e. field_W^1 -> field->get(-1, 0, 0, 1)
:param expr: sympy expression
:param variable_prefix: all variables (and field) are prefixed with this string
this is used for member variables mostly
:param variables_without_prefix: this variables are not prefixed
:return: code string
"""
return cpp_printer.doprint(field_and_symbol_substitute(expr, variable_prefix, variables_without_prefix))
def generate_lattice_model(generation_context, class_name, lb_method, refinement_scaling=None,
**create_kernel_params):
def equations_to_code(equations, variable_prefix="lm.", variables_without_prefix=[], dtype="double"):
def type_eq(eq):
return eq.subs({s: TypedSymbol(s.name, dtype) for s in eq.atoms(sp.Symbol)})
# usually a numpy layout is chosen by default i.e. xyzf - which is bad for waLBerla where at least the spatial
# coordinates should be ordered in reverse direction i.e. zyx
optimization = {'field_layout': 'fzyx'}
if isinstance(equations, AssignmentCollection):
equations = equations.all_assignments
create_kernel_params = default_create_kernel_parameters(generation_context, create_kernel_params)
if create_kernel_params['target'] == 'gpu':
raise ValueError("Lattice Models can only be generated for CPUs. To generate LBM on GPUs use sweeps directly")
variables_without_prefix = list(variables_without_prefix)
c_backend = CBackend()
result = []
left_hand_side_names = [e.lhs.name for e in equations]
for eq in equations:
assignment = SympyAssignment(type_eq(eq.lhs),
field_and_symbol_substitute(eq.rhs, variable_prefix,
variables_without_prefix + left_hand_side_names))
result.append(c_backend(assignment))
return "\n".join(result)
def generate_lattice_model(lattice_model_name=None, optimization={}, refinement_scaling=None,
lb_method=None, **kwargs):
"""
Creates a walberla lattice model consisting of a source and header file
Args:
lattice_model_name: name of the generated lattice model class. If None, it is assumed that this function was
called from a .gen.py file and the filename is taken as lattice model name
optimization: see documentation of create_lb_ast
kwargs: see documentation of create_lb_ast
refinement_scaling: dict from parameter symbol (e.g. relaxation_rate, force parameter) to an expression
how the parameter scales on refined blocks. The refinement factor is represented by
the global symbol REFINEMENT_SCALE_FACTOR
Returns:
tuple with code strings (header, sources)
"""
if lattice_model_name is None:
script_file_name = inspect.stack()[-1][1]
if script_file_name.endswith(".cuda.gen.py"):
raise ValueError("GPU Lattice Model are not yet supported")
elif script_file_name.endswith(".gen.py"):
file_name = script_file_name[:-len(".gen.py")]
lattice_model_name = os.path.split(file_name)[1]
else:
raise ValueError("Not called from a .gen.py file and lattice_model_name is missing")
if 'field_layout' not in optimization:
# usually a numpy layout is chosen by default i.e. xyzf - which is bad for waLBerla where at least the spatial
# coordinates should be ordered in reverse direction i.e. zyx
optimization['field_layout'] = 'fzyx'
params, opt_params = update_with_default_parameters(kwargs.copy(), optimization.copy())
is_float = not generation_context.double_accuracy
params, opt_params = update_with_default_parameters({'lb_method': lb_method}, optimization)
stencil_name = params['stencil']
is_float = not opt_params['double_precision']
if params['force_model'] != 'none' and params['force'] == (0, 0, 0):
params['force'] = sp.symbols("force:3")
stencil_name = get_stencil_name(lb_method.stencil)
if not stencil_name:
raise ValueError("lb_method uses a stencil that is not supported in waLBerla")
params['field_name'] = 'pdfs'
params['temporary_field_name'] = 'pdfs_tmp'
if not lb_method:
lb_method = create_lb_method(**params)
del params['lb_method']
stream_collide_ast = create_lb_ast(lb_method=lb_method, optimization=opt_params, **params)
stream_collide_update_rule = create_lb_update_rule(optimization=opt_params, **params)
stream_collide_ast = create_kernel(stream_collide_update_rule, **create_kernel_params)
stream_collide_ast.function_name = 'kernel_streamCollide'
params['kernel_type'] = 'collide_only'
collide_ast = create_lb_ast(lb_method=lb_method, optimization=opt_params, **params)
collide_update_rule = create_lb_update_rule(optimization=opt_params, **params)
collide_ast = create_kernel(collide_update_rule, **create_kernel_params)
collide_ast.function_name = 'kernel_collide'
dtype = np.float32 if is_float else np.float64
......@@ -152,13 +59,9 @@ def generate_lattice_model(lattice_model_name=None, optimization={}, refinement_
dst_field_name=params['temporary_field_name'],
generic_field_type=dtype,
generic_layout=opt_params['field_layout'])
stream_ast = create_kernel(stream_update_rule.all_assignments)
stream_ast = create_kernel(stream_update_rule, **create_kernel_params)
stream_ast.function_name = 'kernel_stream'
add_openmp(stream_ast, num_threads=opt_params['openmp'])
add_openmp(collide_ast, num_threads=opt_params['openmp'])
add_openmp(stream_collide_ast, num_threads=opt_params['openmp'])
vel_symbols = lb_method.conserved_quantity_computation.first_order_moment_symbols
rho_sym = sp.Symbol("rho")
pdfs_sym = sp.symbols("f_:%d" % (len(lb_method.stencil),))
......@@ -189,8 +92,8 @@ def generate_lattice_model(lattice_model_name=None, optimization={}, refinement_
'momentum_density': momentum_density_symbols})
constant_suffix = "f" if is_float else ""
context = {
'class_name': lattice_model_name,
jinja_context = {
'class_name': class_name,
'stencil_name': stencil_name,
'D': lb_method.dim,
'Q': len(lb_method.stencil),
......@@ -224,23 +127,12 @@ def generate_lattice_model(lattice_model_name=None, optimization={}, refinement_
env = Environment(loader=PackageLoader('lbmpy_walberla'))
add_pystencils_filters_to_jinja_env(env)
header_file = env.get_template('LatticeModel.tmpl.h').render(**context)
cpp_file = env.get_template('LatticeModel.tmpl.cpp').render(**context)
return header_file, cpp_file, context
header = env.get_template('LatticeModel.tmpl.h').render(**jinja_context)
source = env.get_template('LatticeModel.tmpl.cpp').render(**jinja_context)
def generate_lattice_model_files(class_name, *args, **kwargs):
"""
:param kwargs: see documentation of create_lb_ast, additionally
an instance of RefinementScaling can be passed with the 'refinement_scaling' keyword
"""
from pystencils_walberla.cmake_integration import codegen
def generate_lm():
header, sources, _ = generate_lattice_model(class_name, *args, **kwargs)
return header, sources
codegen.register([class_name + ".h", class_name + ".cpp"], generate_lm)
source_extension = "cpp" if create_kernel_params.get("target", "cpu") == "cpu" else "cu"
generation_context.write_file("{}.h".format(class_name), header)
generation_context.write_file("{}.{}".format(class_name, source_extension), source)
class RefinementScaling:
......@@ -287,3 +179,78 @@ class RefinementScaling:
self.add_scaling(p, scaling_rule)
else:
raise ValueError("Invalid value for viscosity_relaxation_rate")
# ------------------------------------------ Internal ------------------------------------------------------------------
def stencil_switch_statement(stencil, values):
template = Template("""
using namespace stencil;
switch( direction ) {
{% for direction_name, value in dir_to_value_dict.items() -%}
case {{direction_name}}: return {{value}};
{% endfor -%}
default:
WALBERLA_ABORT("Invalid Direction");
}
""")
dir_to_value_dict = {offset_to_direction_string(d): cpp_printer.doprint(v) for d, v in zip(stencil, values)}
return template.render(dir_to_value_dict=dir_to_value_dict)
def field_and_symbol_substitute(expr, variable_prefix="lm.", variables_without_prefix=[]):
variables_without_prefix = [v.name if isinstance(v, sp.Symbol) else v for v in variables_without_prefix]
substitutions = {}
for sym in expr.atoms(sp.Symbol):
if isinstance(sym, Field.Access):
fa = sym
prefix = "" if fa.field.name in variables_without_prefix else variable_prefix
if fa.field.index_dimensions == 0:
substitutions[fa] = sp.Symbol("%s%s->get(x,y,z)" % (prefix, fa.field.name))
else:
assert fa.field.index_dimensions == 1, "walberla supports only 0 or 1 index dimensions"
substitutions[fa] = sp.Symbol("%s%s->get(x,y,z,%s)" % (prefix, fa.field.name, fa.index[0]))
else:
if sym.name not in variables_without_prefix:
substitutions[sym] = sp.Symbol(variable_prefix + sym.name)
return expr.subs(substitutions)
def expression_to_code(expr, variable_prefix="lm.", variables_without_prefix=[]):
"""
Takes a sympy expression and creates a C code string from it. Replaces field accesses by
walberla field accesses i.e. field_W^1 -> field->get(-1, 0, 0, 1)
:param expr: sympy expression
:param variable_prefix: all variables (and field) are prefixed with this string
this is used for member variables mostly
:param variables_without_prefix: this variables are not prefixed
:return: code string
"""
return cpp_printer.doprint(field_and_symbol_substitute(expr, variable_prefix, variables_without_prefix))
def equations_to_code(equations, variable_prefix="lm.", variables_without_prefix=[], dtype="double"):
def type_eq(eq):
return eq.subs({s: TypedSymbol(s.name, dtype) for s in eq.atoms(sp.Symbol)})
if isinstance(equations, AssignmentCollection):
equations = equations.all_assignments
variables_without_prefix = list(variables_without_prefix)
c_backend = CBackend()
result = []
left_hand_side_names = [e.lhs.name for e in equations]
for eq in equations:
assignment = SympyAssignment(type_eq(eq.lhs),
field_and_symbol_substitute(eq.rhs, variable_prefix,
variables_without_prefix + left_hand_side_names))
result.append(c_backend(assignment))
return "\n".join(result)
def get_stencil_name(stencil):
for name in ('D2Q9', 'D3Q15', 'D3Q19', 'D3Q27'):
if stencils_have_same_entries(stencil, get_stencil(name, 'walberla')):
return name
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