Commit 569cc1d8 authored by Martin Bauer's avatar Martin Bauer
Browse files

[API] New interface for lattice model generation

- previously a LatticeBoltzmannMethod object was required to generate
  a lattice model for waLBerla
- problem: output fields, smagorinsky, entropic additions are only done
  when coverting the method to a collision rule
- the new interface now takes a collision rule instead of a LB method

to adapt your app have a look at example usage in
walberla/tests/lbm/codegen
parent 00da4cc2
Pipeline #16794 failed with stage
in 2 minutes and 50 seconds
......@@ -101,7 +101,6 @@ public:
static const real_t wInv[{{Q}}];
static const bool compressible = {% if compressible %}true{% else %}false{% endif %};
static const int equilibriumAccuracyOrder = {{equilibrium_order}};
class Sweep
{
......
import warnings
import numpy as np
import sympy as sp
from jinja2 import Environment, PackageLoader, Template
from sympy.tensor import IndexedBase
import pystencils as ps
from lbmpy.creationfunctions import create_lb_update_rule, update_with_default_parameters
from lbmpy.fieldaccess import StreamPullTwoFieldsAccessor, CollideOnlyInplaceAccessor
from lbmpy.fieldaccess import CollideOnlyInplaceAccessor, StreamPullTwoFieldsAccessor
from lbmpy.relaxationrates import relaxation_rate_scaling
from lbmpy.stencils import get_stencil
from lbmpy.updatekernels import create_stream_pull_only_kernel, create_lbm_kernel
import pystencils as ps
from lbmpy.updatekernels import create_lbm_kernel, create_stream_pull_only_kernel
from pystencils import AssignmentCollection, create_kernel
from pystencils.astnodes import SympyAssignment
from pystencils.backends.cbackend import CBackend, CustomSympyPrinter, get_headers
......@@ -24,39 +26,13 @@ cpp_printer = CustomSympyPrinter()
REFINEMENT_SCALE_FACTOR = sp.Symbol("level_scale_factor")
def generate_lattice_model2(generation_context, class_name, collision_rule, refinement_scaling=None,
**create_kernel_params):
# 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
params, opt_params = update_with_default_parameters({}, {'field_layout': 'fzyx'})
is_float = not generation_context.double_accuracy
dtype = np.float32 if is_float else np.float64
lb_method = collision_rule.method
def __lattice_model(generation_context, class_name, lb_method, stream_collide_ast, collide_ast, stream_ast,
refinement_scaling):
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")
q = len(lb_method.stencil)
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")
src_field, dst_field = ps.fields("pdfs({q}), pdfs_tmp({q})".format(q=q), layout='fzyx')
stream_collide_update_rule = create_lbm_kernel(collision_rule, src_field, dst_field, StreamPullTwoFieldsAccessor())
stream_collide_ast = create_kernel(stream_collide_update_rule, **create_kernel_params)
stream_collide_ast.function_name = 'kernel_streamCollide'
collide_update_rule = create_lbm_kernel(collision_rule, src_field, dst_field, CollideOnlyInplaceAccessor())
collide_ast = create_kernel(collide_update_rule, **create_kernel_params)
collide_ast.function_name = 'kernel_collide'
stream_update_rule = create_stream_pull_only_kernel(lb_method.stencil, None, 'pdfs', 'pdfs_tmp', 'fzyx', dtype)
stream_ast = create_kernel(stream_update_rule, **create_kernel_params)
stream_ast.function_name = 'kernel_stream'
is_float = not generation_context.double_accuracy
dtype = np.float32 if is_float else np.float64
vel_symbols = lb_method.conserved_quantity_computation.first_order_moment_symbols
rho_sym = sp.Symbol("rho")
......@@ -99,8 +75,6 @@ def generate_lattice_model2(generation_context, class_name, collision_rule, refi
'weights': ",".join(str(w.evalf()) + constant_suffix for w in lb_method.weights),
'inverse_weights': ",".join(str((1/w).evalf()) + constant_suffix for w in lb_method.weights),
'equilibrium_order': params['equilibrium_order'],
'equilibrium_from_direction': stencil_switch_statement(lb_method.stencil, equilibrium),
'symmetric_equilibrium_from_direction': stencil_switch_statement(lb_method.stencil, symmetric_equilibrium),
'asymmetric_equilibrium_from_direction': stencil_switch_statement(lb_method.stencil, asymmetric_equilibrium),
......@@ -129,124 +103,43 @@ def generate_lattice_model2(generation_context, class_name, collision_rule, refi
header = env.get_template('LatticeModel.tmpl.h').render(**jinja_context)
source = env.get_template('LatticeModel.tmpl.cpp').render(**jinja_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)
generation_context.write_file("{}.cpp".format(class_name), source)
def generate_lattice_model(generation_context, class_name, lb_method, refinement_scaling=None, update_rule_params={},
def generate_lattice_model(generation_context, class_name, collision_rule, refinement_scaling=None,
**create_kernel_params):
# 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'}
is_float = not generation_context.double_accuracy
dtype = np.float32 if is_float else np.float64
lb_method = collision_rule.method
q = len(lb_method.stencil)
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")
is_float = not generation_context.double_accuracy
update_rule_params['lb_method'] = lb_method
params, opt_params = update_with_default_parameters(update_rule_params, optimization)
src_field = ps.Field.create_generic('pdfs', 3, dtype, index_dimensions=1, layout='fzyx', index_shape=(q,))
dst_field = ps.Field.create_generic('pdfs_tmp', 3, dtype, index_dimensions=1, layout='fzyx', index_shape=(q,))
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'
stream_collide_update_rule = create_lb_update_rule(optimization=opt_params, **params)
stream_collide_update_rule = create_lbm_kernel(collision_rule, src_field, dst_field, StreamPullTwoFieldsAccessor())
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_update_rule = create_lb_update_rule(optimization=opt_params, **params)
collide_update_rule = create_lbm_kernel(collision_rule, src_field, dst_field, CollideOnlyInplaceAccessor())
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
stream_update_rule = create_stream_pull_only_kernel(lb_method.stencil, src_field_name=params['field_name'],
dst_field_name=params['temporary_field_name'],
generic_field_type=dtype,
generic_layout=opt_params['field_layout'])
stream_update_rule = create_stream_pull_only_kernel(lb_method.stencil, None, 'pdfs', 'pdfs_tmp', 'fzyx', dtype)
stream_ast = create_kernel(stream_update_rule, **create_kernel_params)
stream_ast.function_name = 'kernel_stream'
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),))
vel_arr_symbols = [IndexedBase(sp.Symbol('u'), shape=(1,))[i] for i in range(len(vel_symbols))]
momentum_density_symbols = [sp.Symbol("md_%d" % (i,)) for i in range(len(vel_symbols))]
equilibrium = lb_method.get_equilibrium()
equilibrium = equilibrium.new_with_substitutions({a: b for a, b in zip(vel_symbols, vel_arr_symbols)})
_, _, equilibrium = add_types(equilibrium.main_assignments, "float32" if is_float else "float64", False)
equilibrium = sp.Matrix([e.rhs for e in equilibrium])
symmetric_equilibrium = get_symmetric_part(equilibrium, vel_arr_symbols)
asymmetric_equilibrium = sp.expand(equilibrium - symmetric_equilibrium)
force_model = lb_method.force_model
macroscopic_velocity_shift = None
if force_model:
if hasattr(force_model, 'macroscopic_velocity_shift'):
macroscopic_velocity_shift = [expression_to_code(e, "lm.", ["rho"])
for e in force_model.macroscopic_velocity_shift(rho_sym)]
cqc = lb_method.conserved_quantity_computation
eq_input_from_input_eqs = cqc.equilibrium_input_equations_from_init_values(sp.Symbol("rho_in"), vel_arr_symbols)
density_velocity_setter_macroscopic_values = equations_to_code(eq_input_from_input_eqs, dtype=dtype,
variables_without_prefix=['rho_in', 'u'])
momentum_density_getter = cqc.output_equations_from_pdfs(pdfs_sym, {'density': rho_sym,
'momentum_density': momentum_density_symbols})
constant_suffix = "f" if is_float else ""
required_headers = get_headers(stream_collide_ast)
jinja_context = {
'class_name': class_name,
'stencil_name': stencil_name,
'D': lb_method.dim,
'Q': len(lb_method.stencil),
'compressible': lb_method.conserved_quantity_computation.compressible,
'weights': ",".join(str(w.evalf()) + constant_suffix for w in lb_method.weights),
'inverse_weights': ",".join(str((1/w).evalf()) + constant_suffix for w in lb_method.weights),
__lattice_model(generation_context, class_name, lb_method, stream_collide_ast, collide_ast, stream_ast,
refinement_scaling)
'equilibrium_order': params['equilibrium_order'],
'equilibrium_from_direction': stencil_switch_statement(lb_method.stencil, equilibrium),
'symmetric_equilibrium_from_direction': stencil_switch_statement(lb_method.stencil, symmetric_equilibrium),
'asymmetric_equilibrium_from_direction': stencil_switch_statement(lb_method.stencil, asymmetric_equilibrium),
'equilibrium': [cpp_printer.doprint(e) for e in equilibrium],
'macroscopic_velocity_shift': macroscopic_velocity_shift,
'density_getters': equations_to_code(cqc.output_equations_from_pdfs(pdfs_sym, {"density": rho_sym}),
variables_without_prefix=[e.name for e in pdfs_sym], dtype=dtype),
'momentum_density_getter': equations_to_code(momentum_density_getter, variables_without_prefix=pdfs_sym,
dtype=dtype),
'density_velocity_setter_macroscopic_values': density_velocity_setter_macroscopic_values,
'refinement_scaling': refinement_scaling,
'stream_collide_kernel': KernelInfo(stream_collide_ast, ['pdfs_tmp'], [('pdfs', 'pdfs_tmp')], []),
'collide_kernel': KernelInfo(collide_ast, [], [], []),
'stream_kernel': KernelInfo(stream_ast, ['pdfs_tmp'], [('pdfs', 'pdfs_tmp')], []),
'target': 'cpu',
'namespace': 'lbm',
'headers': required_headers,
}
env = Environment(loader=PackageLoader('lbmpy_walberla'))
add_pystencils_filters_to_jinja_env(env)
header = env.get_template('LatticeModel.tmpl.h').render(**jinja_context)
source = env.get_template('LatticeModel.tmpl.cpp').render(**jinja_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)
class RefinementScaling:
......
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