diff --git a/src/pystencils_autodiff/framework_integration/astnodes.py b/src/pystencils_autodiff/framework_integration/astnodes.py
index 99f1da8ad8845b4394b2ef9854f441485ecbf5b4..3f60bf879f4c38a8eac9123a425c5246bc7f8a92 100644
--- a/src/pystencils_autodiff/framework_integration/astnodes.py
+++ b/src/pystencils_autodiff/framework_integration/astnodes.py
@@ -210,11 +210,11 @@ class JinjaCppFile(Node):
def __str__(self):
assert self.TEMPLATE, f"Template of {self.__class__} must be set"
render_dict = {k: (self._print(v)
- if not isinstance(v, (pystencils.Field, pystencils.TypedSymbol)) and v is not None
+ if not isinstance(v, (pystencils.Field, pystencils.TypedSymbol, bool)) and v is not None
else v)
if not isinstance(v, Iterable) or isinstance(v, str)
else [(self._print(a)
- if not isinstance(a, (pystencils.Field, pystencils.TypedSymbol)) and a is not None
+ if not isinstance(a, (pystencils.Field, pystencils.TypedSymbol, bool)) and a is not None
else a)
for a in v]
for k, v in self.ast_dict.items()}
diff --git a/src/pystencils_autodiff/framework_integration/printer.py b/src/pystencils_autodiff/framework_integration/printer.py
index afa694eb20186127fe884f098a562f66f54e6b94..2b9a7718b17734ad3ac0dfb6b26ff7d11901488b 100644
--- a/src/pystencils_autodiff/framework_integration/printer.py
+++ b/src/pystencils_autodiff/framework_integration/printer.py
@@ -1,3 +1,5 @@
+import sympy as sp
+
import pystencils.backends.cbackend
from pystencils.kernelparameters import FieldPointerSymbol
@@ -12,14 +14,17 @@ class FrameworkIntegrationPrinter(pystencils.backends.cbackend.CBackend):
"""
def __init__(self):
- super().__init__(sympy_printer=None,
- dialect='c')
+ super().__init__(dialect='c')
def _print(self, node):
from pystencils_autodiff.framework_integration.astnodes import JinjaCppFile
if isinstance(node, JinjaCppFile):
node.printer = self
- return super()._print(node)
+
+ if isinstance(node, sp.Expr):
+ return self.sympy_printer._print(node)
+ else:
+ return super()._print(node)
def _print_WrapperFunction(self, node):
super_result = super()._print_KernelFunction(node)
diff --git a/src/pystencils_autodiff/walberla.py b/src/pystencils_autodiff/walberla.py
index 8716a7e78211c70c2a5031605b9be1c46f0f2d9f..6147c6c88cc04dedd25f1377ff2ade6ed275ddb2 100644
--- a/src/pystencils_autodiff/walberla.py
+++ b/src/pystencils_autodiff/walberla.py
@@ -8,16 +8,44 @@
"""
+from abc import ABC
+from itertools import chain
from os.path import dirname, join
import jinja2
import sympy as sp
+from stringcase import camelcase, pascalcase
+import pystencils
+from pystencils.astnodes import SympyAssignment
from pystencils.data_types import TypedSymbol
from pystencils_autodiff._file_io import read_template_from_file
from pystencils_autodiff.framework_integration.astnodes import JinjaCppFile
+class FieldType(JinjaCppFile):
+
+ TEMPLATE = jinja2.Template("{{ field_type }}")
+
+ def __init__(self, field: pystencils.Field, on_gpu: bool):
+ from pystencils_walberla.jinja_filters import make_field_type, get_field_fsize
+
+ f_size = get_field_fsize(field)
+ field_type = make_field_type(pystencils.data_types.get_base_type(field.dtype), f_size, on_gpu)
+
+ ast_dict = {'on_gpu': on_gpu,
+ 'field_type': field_type
+ }
+ JinjaCppFile.__init__(self, ast_dict)
+
+ @property
+ def headers(self):
+ if self.ast_dict.on_gpu:
+ return ['"field/GhostLayerField.h"']
+ else:
+ return ['"cuda/GPUField.h"']
+
+
class WalberlaModule(JinjaCppFile):
TEMPLATE = read_template_from_file(join(dirname(__file__), 'walberla_main.tmpl.cpp'))
@@ -32,6 +60,9 @@ class WalberlaMain(JinjaCppFile):
int main( int argc, char ** argv )
{
using namespace walberla;
+ using namespace walberla_user;
+ using namespace walberla::pystencils;
+
Environment {{ walberl_env }}( argc, argv );
{{ body | indent(3) }}
@@ -58,12 +89,12 @@ int main( int argc, char ** argv )
return {self.ast_dict.walberl_env}
-class BlockForrestCreation(JinjaCppFile):
- TEMPLATE = jinja2.Template("""auto {{ blocks }} = walberla_user::createBlockForrest(walberlaEnv);""")
- def __init__(self, block_forrest_name):
+class BlockForestCreation(JinjaCppFile):
+ TEMPLATE = jinja2.Template("""auto {{ blocks }} = walberla_user::createBlockForest(walberlaEnv);""")
+ def __init__(self, block_forest_name):
ast_dict = {
- 'blocks': TypedSymbol(block_forrest_name, "auto")
+ 'blocks': TypedSymbol(block_forest_name, "auto")
}
super().__init__(ast_dict)
@@ -73,7 +104,7 @@ class BlockForrestCreation(JinjaCppFile):
return {self.ast_dict.blocks}
-class UniformBlockForrestFromConfig(BlockForrestCreation):
+class UniformBlockforestFromConfig(BlockForestCreation):
TEMPLATE = jinja2.Template(
"""auto {{ blocks }} = blockforest::createUniformBlockGridFromConfig( walberlaEnv.config() );""")
headers = ['"blockforest/Initialization.h"']
@@ -93,13 +124,9 @@ class UniformBlockForrestFromConfig(BlockForrestCreation):
class DefinitionsHeader(JinjaCppFile):
TEMPLATE = read_template_from_file(join(dirname(__file__), 'walberla_user_defintions.tmpl.hpp'))
- def __init__(self, walberla_module):
- self.main_module = walberla_module
- super().__init__({})
-
- @property
- def config_required(self):
- return {"DomainSetup": {"blocks": [1, 1, 1], "cellsPerBlock": [100, 100, 100]}}
+ def __init__(self, lb_model_name, flag_field_type):
+ self.headers = ['<cstdint>', '"lbm/field/PdfField.h"', f'"{lb_model_name}.h"', '"field/FlagField.h"']
+ super().__init__({'lb_model_name': lb_model_name, 'flag_field_type': flag_field_type})
class Using(JinjaCppFile):
@@ -120,7 +147,7 @@ class Using(JinjaCppFile):
class GetParameter(JinjaCppFile):
TEMPLATE = jinja2.Template(
- 'walberlaEnv.config()->getOneBlock("{{ block }}").getParameter<{{ key.dtype }}>("{{ key }}"{% if default %}, {{ default }}{% endif %})' # noqa
+ 'walberlaEnv.config()->getOneBlock("{{ block }}").getParameter<{{ key.dtype }}>("{{ key }}"{% if default %}, static_cast<{{ key.dtype }}>({{ default }}){% endif %})' # noqa
)
def __init__(self, block: str, key, default_value=None):
@@ -138,6 +165,9 @@ class GetParameter(JinjaCppFile):
def __sympy__(self):
return TypedSymbol(self.ast_dict.key.name + "Config", str(self.ast_dict.key.dtype))
+ def __getattr__(self, name):
+ return getattr(self.__sympy__(), name)
+
class FieldAllocation(JinjaCppFile):
"""
@@ -156,25 +186,40 @@ class FieldAllocation(JinjaCppFile):
""" # noqa
- TEMPLATE = jinja2.Template("""BlockDataID {{ field_name }}_data = field::addToStorage<{{ field_type }}>( {{ block_forrest }},
- {{ field_name }}
- {%- if init_value -%} , {{ init_value }} {% endif %}
- {%- if layout_str -%} , {{ layout_str }} {% endif %}
+ TEMPLATE = jinja2.Template("""
+{%- if on_gpu -%}
+BlockDataID {{ field_name }}_data_gpu = cuda::addGPUFieldToStorage<{{ field_type }}>({{ block_forest }},
+ "{{ field_name }}",
+ {{ f_size }},
+ {{ layout }},
+ {{ num_ghost_layers }},
+ {{ usePitchedMem }} );
+{%- else -%}
+BlockDataID {{ field_name }}_data = field::addToStorage<{{ field_type }}>( {{ block_forest }},
+ "{{ field_name }}"
+ {%- if init_value -%} , {{ init_value }}{% endif %}
+ {%- if layout -%} , {{ layout }}{% endif %}
{%- if num_ghost_layers -%}, {{ num_ghost_layers }} {% endif %}
{%- if always_initialize -%}, {{ always_initialize }} {% endif %});
+{%- endif %}
""") # noqa
- def __init__(self, block_forrest, field):
- self._symbol = TypedSymbol(field.name + '_data', 'BlockDataID')
+ def __init__(self, block_forest, field, on_gpu=False, usePitchedMem=True, num_ghost_layers=1):
+ self._symbol = TypedSymbol(field.name + ('_data_gpu' if on_gpu else '_data'), 'BlockDataID')
ast_dict = {
- 'block_forrest': block_forrest,
+ 'block_forest': block_forest,
'field_name': field.name,
- 'field_type': f'GhostLayerField< {field.dtype}, {field.index_shape[0] if field.index_shape else 1} >'
+ # f'GhostLayerField< {field.dtype}, {field.index_shape[0] if field.index_shape else 1} >'
+ 'field_type': FieldType(field, on_gpu),
+ 'on_gpu': on_gpu,
+ 'f_size': field.index_shape[0] if field.index_shape else 1,
+ 'init_value': f'{field.dtype}{{}}',
+ 'num_ghost_layers': num_ghost_layers,
+ 'layout': 'field::zyxf',
+ 'usePitchedMem': 'true' if usePitchedMem else 'false',
}
super().__init__(ast_dict)
- headers = ['"field/GhostLayerField.h"']
-
@property
def symbol(self):
return self._symbol
@@ -183,6 +228,12 @@ class FieldAllocation(JinjaCppFile):
def symbols_defined(self):
return {self.symbol}
+ @property
+ def headers(self):
+ return (['"cuda/AddGPUFieldToStorage.h"', '"field/GhostLayerField.h"']
+ if self.ast_dict.on_gpu
+ else ['"field/AddToStorage.h"'])
+
class WalberlaVector(JinjaCppFile):
TEMPLATE = jinja2.Template("""math::Vector{{ndim}}<{{dtype}}>({{offsets}})""")
@@ -212,7 +263,7 @@ class PdfFieldAllocation(FieldAllocation):
const Set < SUID > & requiredSelectors = Set < SUID > : : emptySet(),
const Set < SUID > & incompatibleSelectors = Set < SUID > : : emptySet() )
"""
- TEMPLATE = jinja2.Template("""BlockDataID {{field_name}}_data = lbm::field::addPdfFieldToStorage < {{ field_type }} > ({{ block_forrest }},
+ TEMPLATE = jinja2.Template("""BlockDataID {{field_name}}_data = lbm::field::addPdfFieldToStorage < {{ field_type }} > ({{ block_forest }},
{{field_name}},
{{lattice_model}}
{%- if initial_velocity -%} , {{initial_velocity }} {% endif %}
@@ -220,8 +271,8 @@ class PdfFieldAllocation(FieldAllocation):
{%- if num_ghost_layers -%}, {{num_ghost_layers }} {% endif %});
""") # noqa
- def __init__(self, block_forrest, field, lattice_model, initial_velocity=None, initial_density=None):
- super().__init__(block_forrest, field)
+ def __init__(self, block_forest, field, lattice_model, initial_velocity=None, initial_density=None, on_gpu=False):
+ super().__init__(block_forest, field, on_gpu)
if initial_velocity and not isinstance(initial_velocity, WalberlaVector):
initial_velocity = WalberlaVector(initial_velocity)
@@ -231,7 +282,7 @@ class PdfFieldAllocation(FieldAllocation):
'initial_velocity': initial_velocity,
})
- headers = ['lbm/field/AddToStorage.h']
+ headers = ['"lbm/field/AddToStorage.h"']
class FlagFieldAllocation(FieldAllocation):
@@ -248,12 +299,480 @@ class FlagFieldAllocation(FieldAllocation):
const Set<SUID> & requiredSelectors = Set<SUID>::emptySet(),
const Set<SUID> & incompatibleSelectors = Set<SUID>::emptySet() )
""" # noqa
- TEMPLATE = jinja2.Template("""BlockDataID {{field_name}}_data = field::addFlagFieldToStorage < {{ field_type }} > ({{ block_forrest }},
+ TEMPLATE = jinja2.Template("""BlockDataID {{field_name}}_data = field::addFlagFieldToStorage < {{ field_type }} > ({{ block_forest }},
{{field_name}}
{%- if num_ghost_layers -%}, {{num_ghost_layers }} {% endif %});
""") # noqa
- def __init__(self, block_forrest, field):
- super().__init__(block_forrest, field)
+ def __init__(self, block_forest, field, on_gpu=False):
+ super().__init__(block_forest, field, on_gpu)
+
+ headers = ['"field/AddToStorage.h"']
+
+
+class FlagUidDefinition(JinjaCppFile):
+ TEMPLATE = jinja2.Template('const FlagUID {{ name }}FlagUID("{{ name }}");')
+
+ def __init__(self, name):
+ self._symbol = TypedSymbol(name + 'FlagUID', 'FlagUID')
+ ast_dict = {
+ 'name': name,
+ }
+ super().__init__(ast_dict)
+
+ headers = ['"field/FlagUID.h"']
+
+ @property
+ def symbol(self):
+ return self._symbol
+
+ @property
+ def symbols_defined(self):
+ return {self.symbol}
+
+
+class BoundaryHandling(ABC):
+ pass
+
+
+class BoundaryHandlingFromConfig(JinjaCppFile):
+
+ TEMPLATE = jinja2.Template("""auto {{ boundaries_config }} = walberlaEnv.config()->getOneBlock( "Boundaries" );
+geometry::initBoundaryHandling<FlagField_T>(*{{ block_forest }}, {{ flag_field_id }}, {{ boundaries_config }});
+geometry::setNonBoundaryCellsToDomain<FlagField_T>(*{{ block_forest }}, {{ flag_field_id }}, {{ fluid_uid }});""")
+
+ def __init__(self, block_forest, flag_field_id, fluid_uid):
+ self._symbol = boundaries_config = TypedSymbol('boundariesConfig', 'auto')
+ ast_dict = {
+ 'flag_field_id': flag_field_id,
+ 'block_forest': block_forest,
+ 'fluid_uid': fluid_uid,
+ 'boundaries_config': boundaries_config,
+ }
+ super().__init__(ast_dict)
+
+ @property
+ def symbol(self):
+ return self._symbol
+
+ @property
+ def symbols_defined(self):
+ return {self.symbol}
+
+
+class FillFromFlagField(JinjaCppFile):
+
+ TEMPLATE = jinja2.Template("""
+ {{ boundary_condition }}.fillFromFlagField<{{ flag_field_type }}>( {{ block_forest }}, {{ flag_field_id }}, FlagUID("{{ boundary_name }}"), {{ flag_field_id }});
+ """) # noqa
+
+ def __init__(self, flag_field_id, fluid_uid):
+ self._symbol = boundaries_config = TypedSymbol('boundariesConfig', 'auto')
+ ast_dict = {
+ 'flag_field_id': flag_field_id,
+ 'fluid_uid': fluid_uid,
+ 'boundaries_config': boundaries_config,
+ }
+ super().__init__(ast_dict)
+
+ @property
+ def symbol(self):
+ return self._symbol
+
+ @property
+ def symbols_defined(self):
+ return {self.symbol}
+
+
+class LbCommunicationSetup(JinjaCppFile):
+
+ TEMPLATE = jinja2.Template("""blockforest::communication::UniformBufferedScheme<lbm::{{ lb_model_type }}::CommunicationStencil> {{ communication }}( blocks );
+{{ communication }}.addPackInfo( make_shared< lbm::PdfFieldPackInfo<lbm::{{ lb_model_type }}> >( {{ pdf_id }} ) );
+ """) # noqa
+
+ def __init__(self, lb_model_type, pdf_id):
+ self._symbol = TypedSymbol('communication', 'auto')
+ ast_dict = {
+ 'lb_model_type': lb_model_type,
+ 'pdf_id': pdf_id,
+ 'communication': self._symbol,
+ }
+ super().__init__(ast_dict)
+
+ @property
+ def symbol(self):
+ return self._symbol
+
+ @property
+ def symbols_defined(self):
+ return {self.symbol}
+
+ headers = ['"blockforest/communication/UniformBufferedScheme.h"', '"lbm/communication/PdfFieldPackInfo.h"']
+
+
+class BeforeFunction(JinjaCppFile):
+
+ TEMPLATE = jinja2.Template("""BeforeFunction( {{ f }}, "{{ f }}" )""") # noqa
+
+ def __init__(self, function):
+ ast_dict = {
+ 'f': function,
+ }
+ super().__init__(ast_dict)
+
+
+class AfterFunction(JinjaCppFile):
+
+ TEMPLATE = jinja2.Template("""AfterFunction( {{ f }}, "{{ f }}" )""") # noqa
+
+ def __init__(self, function):
+ ast_dict = {
+ 'f': function,
+ }
+ super().__init__(ast_dict)
+
+
+class Sweep(JinjaCppFile):
+
+ TEMPLATE = jinja2.Template("""Sweep( {{ f }}, "{{ f }}" )""") # noqa
+
+ def __init__(self, function):
+ ast_dict = {
+ 'f': function,
+ }
+ super().__init__(ast_dict)
+
+
+class TimeLoop(JinjaCppFile):
+
+ TEMPLATE = jinja2.Template("""
+{{ timeloop_definition }}
+
+{% for f in before_functions -%}
+{{timeloop_name}}.add() << {{ f }};
+{% endfor %}
+{% for s in sweeps -%}
+{{timeloop_name}}.add() << {{ s }};
+{% endfor %}
+{% for f in after_functions -%}
+{{timeloop_name}}.add() << {{ f }};
+{% endfor %}
+ """) # noqa
+
+ def __init__(self, block_forest, before_functions, sweeps, after_functions, timesteps, timeloop_name='timeloop'):
+ self._symbol = TypedSymbol(timeloop_name, 'SweepTimeloop')
+ ast_dict = {
+ 'timeloop_definition': CppObjectConstruction(self._symbol,
+ [f'{block_forest}->getBlockStorage()',
+ timesteps]),
+ 'before_functions': [BeforeFunction(f) for f in before_functions],
+ 'sweeps': [Sweep(f) for f in sweeps],
+ 'after_functions': [AfterFunction(f) for f in after_functions],
+ 'timeloop_name': timeloop_name
+ }
+ super().__init__(ast_dict)
+
+ @property
+ def symbol(self):
+ return self._symbol
+
+ @property
+ def symbols_defined(self):
+ return {self.symbol}
+
+ headers = ['"timeloop/all.h"']
+
+
+class U_Rho_Adaptor(JinjaCppFile):
+ """Docstring for U_Rho_Adaptor. """
+
+ TEMPLATE = jinja2.Template("""
+field::addFieldAdaptor<lbm::Adaptor<lbm::{{ lb_model_type }}>::Density> ( {{ block_forest }}, {{ pdf_id }}, "{{ density_name }}" );
+field::addFieldAdaptor<lbm::Adaptor<lbm::{{ lb_model_type }}>::VelocityVector>( {{ block_forest }}, {{ pdf_id }}, "{{ velocity_name }}" );
+ """) # noqa
+
+ def __init__(self, block_forest, lb_model_type, pdf_id):
+ self.density_name = "DensityAdaptor"
+ self.velocity_name = "VelocityAdaptor"
+ ast_dict = {
+ 'lb_model_type': lb_model_type,
+ 'pdf_id': pdf_id,
+ 'block_forest': block_forest,
+ 'density_name': self.density_name,
+ 'velocity_name': self.velocity_name,
+ }
+ super().__init__(ast_dict)
+
+ headers = ['"timeloop/all.h"']
+
+
+class RunTimeLoop(JinjaCppFile):
+ TEMPLATE = jinja2.Template("""{%- if with_gui == 'true' %}
+if( parameters.getParameter<bool>( "useGui", {{ use_gui_default }}) )
+{
+ GUI gui ( {{ timeloop }}, {{ block_forest }}, argc, argv );
+ lbm::connectToGui<LatticeModel_T> ( gui );
+ gui.run();
+}
+else {
+ {{ timeloop }}.run();
+}
+{%- else %}
+timeloop.run();
+{%- endif %}
+ """) # noqa
+
+ def __init__(self, block_forest, timeloop, with_gui=False, use_gui_default='false'):
+ self.density_name = "DensityAdaptor"
+ self.velocity_name = "VelocityAdaptor"
+ ast_dict = {
+ 'block_forest': block_forest,
+ 'timeloop': timeloop.symbol,
+ 'with_gui': with_gui,
+ 'use_gui_default': use_gui_default,
+ }
+ super().__init__(ast_dict)
+
+ @property
+ def headers(self):
+ if self.ast_dict.with_gui:
+ return [
+ "gui/all.h",
+ "lbm/gui/Connection.h"
+ ]
+ else:
+ return []
+
+
+class CppObjectConstruction(JinjaCppFile, pystencils.astnodes.SympyAssignment):
+
+ TEMPLATE = jinja2.Template("""{{ symbol.dtype }} {{ symbol }}({{ arg_list }});""") # noqa
+
+ def __init__(self, symbol, args):
+ JinjaCppFile.__init__(self, {})
+ self.ast_dict.update({
+ 'symbol': symbol,
+ 'args': args,
+ 'arg_list': ', '.join(self.printer(a) for a in args)
+ })
+ pystencils.astnodes.SympyAssignment.__init__(self, symbol, 1)
+
+
+class ResolveUndefinedSymbols(JinjaCppFile):
+
+ TEMPLATE = jinja2.Template("""
+{% for p in parameters %}
+{{ p }}
+{% endfor %}
+{{ block }}
+""") # noqa
+
+ def __init__(self, block, config_block):
+ self.block = block
+ JinjaCppFile.__init__(self, {})
+ self.ast_dict.update({
+ 'block': block,
+ 'config_block': config_block,
+ 'parameters': self.parameter_definitions
+ })
+
+ @property
+ def symbols_defined(self):
+ return self.block.undefined_symbols
+
+ def __repr__(self):
+ self.ast_dict.parameters = self.parameter_definitions
+ return super().__repr__()
+
+ def __str__(self):
+ self.ast_dict.parameters = self.parameter_definitions
+ return super().__str__()
+
+ @property
+ def parameter_definitions(self):
+ parameter_definitions = [SympyAssignment(s, GetParameter(self.ast_dict.config_block, s))
+ for s in self.block.undefined_symbols]
+ return parameter_definitions
+
+ @property
+ def config_required(self):
+ return {self.ast_dict.config_block, {s.name: None for s in self.symbols_defined}}
+
+
+class FieldCopyFunctor(JinjaCppFile):
+
+ TEMPLATE = jinja2.Template("""cuda::fieldCpyFunctor<{{from_type}}, {{ to_type }} >({{from_id}}, {{to_id}})""") # noqa
+
+ def __init__(self, from_id, from_type, to_id, to_type):
+ ast_dict = {'from_id': from_id,
+ 'from_type': from_type,
+ 'to_id': to_id,
+ 'to_type': to_type
+ }
+ super().__init__(ast_dict)
+
+ headers = ['"cuda/FieldCopy.h"']
+
+
+class AllocateAllFields(JinjaCppFile):
+
+ TEMPLATE = jinja2.Template("""
+// CPU Arrays
+{% for c in cpu_arrays %}
+{{- c }}
+{% endfor %}
+// GPU Arrays
+{% for g in gpu_arrays %}
+{{- g }}
+{% endfor %}
+{{ block }}
+""") # noqa
+
+ def __init__(self, block_forest, data_handling, lb_index_shape=(-1,), lb_model_name=None):
+ self.data_handling = data_handling
+ self.block_forest = block_forest
+ self.lb_index_shape = lb_index_shape
+ self.lb_model_name = lb_model_name
+ JinjaCppFile.__init__(self, {})
+ self._update()
+
+ def __repr__(self):
+ self._update()
+ return super().__repr__()
+
+ def __str__(self):
+ self._update()
+ return super().__str__()
+
+ def _update(self):
+ self.ast_dict.update({
+ 'cpu_arrays': self._cpu_allocations.values(),
+ 'gpu_arrays': self._gpu_allocations.values()
+ })
+
+ @property
+ def _cpu_allocations(self):
+ return {s: FieldAllocation(self.block_forest, self.data_handling.fields[s])
+ for s in self.data_handling.cpu_arrays.keys()}
+ # if not (self.data_handling.fields[s].index_shape == self.lb_index_shape)}
+
+ @property
+ def _gpu_allocations(self):
+ return {s: FieldAllocation(self.block_forest, self.data_handling.fields[s], on_gpu=True)
+ for s in self.data_handling.gpu_arrays.keys()}
+ # if not (self.data_handling.fields[s].index_shape == self.lb_index_shape)}
+
+
+class InitBoundaryHandling(JinjaCppFile):
+ TEMPLATE = jinja2.Template("""
+// This is my fluid 🏊. It's special... 🤔 but for the computer just a 0
+{{ fluid_uid_definition }}
+
+// Initialize geometry 🎲
+{{ geometry_initialization }}
+
+// Here are the generated boundaries. They are not so special... 👎
+{%- for b in generated_boundaries %}
+{{ b }}
+{% endfor %}
+""") # noqa
+
+ def __init__(self, block_forest, flag_field_id, pdf_field_id, boundary_conditions):
+ self.fluid = FlagUidDefinition("fluid")
+ ast_dict = {'fluid_uid_definition': self.fluid,
+ 'geometry_initialization': BoundaryHandlingFromConfig(block_forest,
+ flag_field_id,
+ self.fluid.symbol),
+ 'generated_boundaries': [GeneratedBoundaryInitialization(block_forest,
+ b,
+ pdf_field_id,
+ flag_field_id,
+ self.fluid.symbol)
+ for b in boundary_conditions]
+ }
+ super().__init__(ast_dict)
+
+ headers = ['"cuda/FieldCopy.h"', '"geometry/InitBoundaryHandling.h"']
+
+
+class GeneratedBoundaryInitialization(JinjaCppFile):
+ TEMPLATE = jinja2.Template("""lbm::{{ boundary_condition }} {{ identifier }}( {{ block_forest }}, {{ pdf_field_id }} );
+{{ identifier }}.fillFromFlagField<FlagField_T>( {{ block_forest }}, {{ flag_field_id }}, FlagUID("{{ boundary_condition }}"), {{ fluid_uid }} );
+""") # noqa
+
+ def __init__(self, block_forest, boundary_condition, pdf_field_id, flag_field_id, fluid_uid):
+ self.fluid = FlagUidDefinition("fluid")
+ ast_dict = {'block_forest': block_forest,
+ 'boundary_condition': pascalcase(boundary_condition.name),
+ 'identifier': camelcase(boundary_condition.name),
+ 'pdf_field_id': pdf_field_id,
+ 'fluid_uid': fluid_uid,
+ 'flag_field_id': flag_field_id,
+ }
+ super().__init__(ast_dict)
+
+ @property
+ def headers(self):
+ return [f'"{pascalcase(self.ast_dict.boundary_condition)}.h"']
+
+ @property
+ def symbols_defined(self):
+ # Could also be defined
+ # TypedSymbol(self.ast_dict.identifier, 'FlagUID'),
+ return {TypedSymbol(self.ast_dict.identifier, f'lbm::{self.ast_dict.boundary_condition}')}
+
+
+class SweepCreation(JinjaCppFile):
+ TEMPLATE = jinja2.Template("""{{ sweep_class_name }}( {{ parameter_str }} )""") # noqa
+
+ def __init__(self, sweep_class_name: str, field_allocation: AllocateAllFields, ast, parameters_to_ignore=None):
+ parameters = ast.get_parameters()
+ parameter_ids = [field_allocation._cpu_allocations[p.symbol.name.replace('_data_', '')].symbol.name
+ if ast.target == 'cpu'
+ else field_allocation._gpu_allocations[p.symbol.name.replace('_data_', '')].symbol.name
+ for p in parameters
+ if p.is_field_pointer or not p.is_field_parameter
+ ]
+
+ ast_dict = {'sweep_class_name': sweep_class_name,
+ 'parameter_ids': parameter_ids,
+ 'parameter_str': ', '.join(parameter_ids)}
+ super().__init__(ast_dict)
+
+ @property
+ def headers(self):
+ return [f'"{self.ast_dict.sweep_class_name}.h"']
+
+
+class SweepOverAllBlocks(JinjaCppFile):
+ # TEMPLATE = jinja2.Template("""std::for_each({{block_forest}}->begin(), {{block_forest}}->end(), {{functor}});""") # noqa
+ TEMPLATE = jinja2.Template("""auto {{sweep_class_name | lower() }} = {{functor}};
+for( auto& block : *{{block_forest}}) {{sweep_class_name | lower() }}(&block);""") # noqa
+
+ def __init__(self, functor: SweepCreation, block_forest):
+ ast_dict = {'functor': functor,
+ 'sweep_class_name': functor.ast_dict.sweep_class_name,
+ 'block_forest': block_forest}
+ super().__init__(ast_dict)
+
+ headers = ['<algorithm>']
+
+
+# class GetSweep(JinjaCppFile):
+ # TEMPLATE = jinja2.Template("""std::function<void(*domain_decomposition::IBlock)>({{ sweep_class }})""") # noqa
+
+ # def __init__(self, sweep_class: SweepCreation):
+
+ # ast_dict = {'sweep_class': sweep_class, }
+ # super().__init__(ast_dict)
+
+ # headers = ['"domain_decomposition/IBlock.h"']
+
+
+# class MakeSharedFunctor(JinjaCppFile):
+ # TEMPLATE = jinja2.Template("""makeSharedSweep{{ functor }}(&{{block_forest}})""") # noqa
- headers = ['field/AddToStorage.h']
+ # def __init__(self, functor, block_forest):
+ # ast_dict = {'functor': functor,
+ # 'block_forest': block_forest}
+ # super().__init__(ast_dict)
diff --git a/src/pystencils_autodiff/walberla_user_defintions.tmpl.hpp b/src/pystencils_autodiff/walberla_user_defintions.tmpl.hpp
index dc6d4e5d2f623b60b92a86fb8df4682198152e12..0df3276c1d924ade412a667fe6a39bb18e3d53d8 100644
--- a/src/pystencils_autodiff/walberla_user_defintions.tmpl.hpp
+++ b/src/pystencils_autodiff/walberla_user_defintions.tmpl.hpp
@@ -7,22 +7,21 @@
#pragma once
-#include "lbm/communication/PdfFieldPackInfo.h"
-#include "lbm/field/AddToStorage.h"
-#include "lbm/field/PdfField.h"
-#include "lbm/gui/Connection.h"
-#include "lbm/vtk/VTKOutput.h"
+
+{% for header in headers -%}
+#include {{ header }}
+{% endfor %}
namespace walberla_user {
using namespace walberla;
-using LatticeModel_T = lbm::LbCodeGenerationExample_LatticeModel;
+using LatticeModel_T = lbm::{{ lb_model_name }};
using Stencil_T = LatticeModel_T::Stencil;
-using CommunicationStencil_T = LatticeModel_T::CommunicationStencil_T;
-using lPdfField_T = bm::PdfField<LatticeModel_T>;
-using flag_t = walberla::uint8_t;
-using FlagField_T FlagField<flag_t>;
+using CommunicationStencil_T = LatticeModel_T::CommunicationStencil;
+using PdfField_T = lbm::PdfField<LatticeModel_T>;
+using flag_t = {{ flag_field_type }};
+using FlagField_T = FlagField<flag_t>;
using PdfField_T = lbm::PdfField<LatticeModel_T>;
-typedef VectorField_T = GhostLayerField<real_t, LatticeModel_T::Stencil::D>;
-typedef ScalarField_T GhostLayerField<real_t, 1>;
+using VectorField_T = GhostLayerField<real_t, LatticeModel_T::Stencil::D>;
+using ScalarField_T = GhostLayerField<real_t, 1>;
} // namespace walberla_user
diff --git a/src/pystencils_autodiff/wald_und_wiesen_simulation.py b/src/pystencils_autodiff/wald_und_wiesen_simulation.py
new file mode 100644
index 0000000000000000000000000000000000000000..ee142fafe3af703777b17323d7e9998901c5dd98
--- /dev/null
+++ b/src/pystencils_autodiff/wald_und_wiesen_simulation.py
@@ -0,0 +1,141 @@
+#
+# Copyright © 2020 Stephan Seitz <stephan.seitz@fau.de>
+#
+# Distributed under terms of the GPLv3 license.
+
+"""
+
+"""
+
+from typing import Dict
+
+import sympy as sp
+from stringcase import pascalcase
+
+import lbmpy_walberla
+import pystencils
+import pystencils_walberla.codegen
+from pystencils.astnodes import Block, EmptyLine
+from pystencils_autodiff.walberla import (
+ AllocateAllFields, DefinitionsHeader, InitBoundaryHandling, LbCommunicationSetup,
+ ResolveUndefinedSymbols, RunTimeLoop, SweepCreation, SweepOverAllBlocks, TimeLoop,
+ UniformBlockforestFromConfig, WalberlaMain, WalberlaModule)
+
+
+class WaldUndWiesenSimulation():
+ def _get_sweep_class_name(prefix='Kernel'):
+
+ ctr = 0
+ while True:
+ yield f'{prefix}{ctr}'
+ ctr += 1
+
+ def __init__(self,
+ graph_data_handling,
+ codegen_context,
+ boundary_handling: pystencils.boundaries.BoundaryHandling = None,
+ lb_rule=None,
+ refinement_scaling=None):
+ self._data_handling = graph_data_handling
+ self._lb_rule = lb_rule
+ self._refinement_scaling = refinement_scaling
+ self._block_forest = UniformBlockforestFromConfig()
+ self.parameter_config_block = 'parameters'
+ self._codegen_context = codegen_context
+ self._boundary_handling = boundary_handling
+ self._lb_model_name = 'GeneratedLatticeModel'
+ self._flag_field_dtype = 'uint32_t'
+ self._kernel_class_generator = WaldUndWiesenSimulation._get_sweep_class_name()
+ self._with_gui = False
+ self._with_gui_default = False
+
+ def _create_helper_files(self) -> Dict[str, str]:
+ if self._lb_rule:
+ lbmpy_walberla.generate_lattice_model(self._codegen_context, self._lb_model_name,
+ self._lb_rule,
+ refinement_scaling=self._refinement_scaling)
+ if self._boundary_handling:
+ for bc in self.boundary_conditions:
+ lbmpy_walberla.generate_boundary(
+ self._codegen_context, pascalcase(bc.name), bc, self._lb_rule.method)
+
+ def _create_module(self):
+ if self._lb_rule:
+ lb_shape = (len(self._lb_rule.method.stencil),)
+ else:
+ lb_shape = (-1,)
+
+ self._field_allocations = field_allocations = AllocateAllFields(self._block_forest.blocks,
+ self._data_handling,
+ lb_shape,
+ self._lb_model_name)
+
+ if self._boundary_handling:
+ flag_field_id = field_allocations._cpu_allocations[
+ self._boundary_handling.flag_interface.flag_field_name].symbol
+
+ if self._lb_rule:
+ pdf_field_id = field_allocations._gpu_allocations.get(
+ 'ldc_pdfSrc', field_allocations._cpu_allocations['ldc_pdfSrc']).symbol
+ else:
+ pdf_field_id = None
+
+ call_nodes = filter(lambda x: x, [self._graph_to_sweep(c) for c in self._data_handling.call_queue])
+
+ module = WalberlaModule(WalberlaMain(Block([
+ self._block_forest,
+ ResolveUndefinedSymbols(
+ Block([
+ field_allocations,
+ InitBoundaryHandling(self._block_forest.blocks, flag_field_id,
+ pdf_field_id, self.boundary_conditions)
+ if self._boundary_handling else EmptyLine(),
+ LbCommunicationSetup(self._lb_model_name,
+ pdf_field_id)
+ if self._lb_rule else EmptyLine(),
+ *call_nodes
+ ]), self.parameter_config_block)
+ ])))
+
+ self._codegen_context.write_file("main.cpp", str(module))
+ return module
+
+ def _create_defintions_header(self):
+ self._codegen_context.write_file("UserDefinitions.h",
+ str(DefinitionsHeader(self._lb_model_name, self._flag_field_dtype)))
+
+ def write_files(self):
+ self._create_helper_files()
+ self._create_module()
+ self._create_defintions_header()
+
+ @property
+ def boundary_conditions(self):
+ return self._boundary_handling._boundary_object_to_boundary_info.keys()
+
+ def _graph_to_sweep(self, c):
+ from pystencils_autodiff.graph_datahandling import KernelCall, TimeloopRun
+
+ if isinstance(c, KernelCall):
+ sweep_class_name = next(self._kernel_class_generator)
+ pystencils_walberla.codegen.generate_sweep(
+ self._codegen_context, sweep_class_name, c.kernel.ast)
+ rtn = SweepOverAllBlocks(SweepCreation(sweep_class_name, self._field_allocations,
+ c.kernel.ast), self._block_forest.blocks)
+
+ elif isinstance(c, TimeloopRun):
+ sweeps = []
+ for a in c.timeloop._single_step_asts:
+ if 'indexField' in [f.name for f in a.fields_accessed]:
+ continue
+ sweep_class_name = next(self._kernel_class_generator)
+ pystencils_walberla.codegen.generate_sweep(
+ self._codegen_context, sweep_class_name, a)
+ sweeps.append(SweepCreation(sweep_class_name, self._field_allocations, a))
+
+ loop = TimeLoop(self._block_forest.blocks, [], sweeps, [], sp.S(c.time_steps))
+ rtn = Block([loop, RunTimeLoop(self._block_forest.blocks, loop, self._with_gui, self._with_gui_default)])
+
+ else:
+ return None
+ return rtn
diff --git a/tests/test_walberla.py b/tests/test_walberla.py
index 5f2859e6e8327d04b6a1dc2bb0bc8e8005bf86af..162dea9bcf6437702d13f18aa206a20e70b061f5 100644
--- a/tests/test_walberla.py
+++ b/tests/test_walberla.py
@@ -6,16 +6,24 @@
"""
"""
+import os
+import sys
+from os.path import dirname, join
+
import numpy as np
-import sympy as sp
+import lbmpy_walberla
import pystencils
+from lbmpy.creationfunctions import create_lb_collision_rule, create_lbm_kernel
from pystencils.astnodes import Block, EmptyLine, SympyAssignment
-from pystencils.data_types import TypedSymbol
+from pystencils.data_types import TypedSymbol, create_type
from pystencils_autodiff._file_io import write_file
+from pystencils_autodiff.graph_datahandling import GraphDataHandling
from pystencils_autodiff.walberla import (
- DefinitionsHeader, FieldAllocation, GetParameter, PdfFieldAllocation,FlagFieldAllocation,
- UniformBlockForrestFromConfig, WalberlaMain, WalberlaModule)
+ DefinitionsHeader, FieldAllocation, FlagFieldAllocation, GetParameter, PdfFieldAllocation,
+ ResolveUndefinedSymbols, UniformBlockforestFromConfig, WalberlaMain, WalberlaModule)
+from pystencils_autodiff.wald_und_wiesen_simulation import WaldUndWiesenSimulation
+from pystencils_walberla.cmake_integration import ManualCodeGenerationContext
def test_walberla():
@@ -28,18 +36,18 @@ def test_walberla():
number_symbol = TypedSymbol('number', np.float32)
crazy_plus_one = TypedSymbol('crazy', np.float32)
- block_forrest = UniformBlockForrestFromConfig()
+ block_forest = UniformBlockforestFromConfig()
block = Block([
- block_forrest,
+ block_forest,
SympyAssignment(foo_symbol, GetParameter('parameters', foo_symbol)),
SympyAssignment(number_symbol, GetParameter('parameters', number_symbol, 1.2)),
SympyAssignment(crazy_plus_one, number_symbol + 1),
EmptyLine(),
- FieldAllocation(block_forrest.blocks, x),
- PdfFieldAllocation(block_forrest.blocks, pdf, 'LbModel_T'),
- PdfFieldAllocation(block_forrest.blocks, pdf2, 'LbModel_T', [0, 0, 0], 1),
- FlagFieldAllocation(block_forrest.blocks, flags)
+ FieldAllocation(block_forest.blocks, x, on_gpu=False),
+ PdfFieldAllocation(block_forest.blocks, pdf, 'LbModel_T', on_gpu=True),
+ PdfFieldAllocation(block_forest.blocks, pdf2, 'LbModel_T', [0, 0, 0], 1, on_gpu=True),
+ FlagFieldAllocation(block_forest.blocks, flags)
])
module = WalberlaModule(WalberlaMain(block))
@@ -48,5 +56,56 @@ def test_walberla():
write_file('/localhome/seitz_local/projects/walberla/apps/autogen/main.cpp', code)
- definitions = DefinitionsHeader(module)
+ definitions = DefinitionsHeader(module, 'uint8_t')
write_file('/localhome/seitz_local/projects/walberla/apps/autogen/UserDefinitions.h', str(definitions))
+
+
+def test_wald_wiesen_simulation():
+ with ManualCodeGenerationContext() as ctx:
+ dh = GraphDataHandling((30, 30),
+ periodicity=False,
+ default_ghost_layers=1,
+ default_target='cpu')
+ dh.add_arrays('x, y')
+ dh.add_arrays('w, z', gpu=True)
+
+ sim = WaldUndWiesenSimulation(dh, ctx)
+ print(sim._create_module())
+
+
+def test_wald_wiesen_lbm():
+ sys.path.append(dirname(__file__))
+ with ManualCodeGenerationContext() as ctx:
+ from test_graph_datahandling import ldc_setup
+ opt_params = {'target': 'gpu'}
+ lbm_step = ldc_setup(domain_size=(30, 30), optimization=opt_params, fixed_loop_sizes=False)
+
+ sim = WaldUndWiesenSimulation(lbm_step.data_handling,
+ ctx,
+ lbm_step.boundary_handling,
+ create_lb_collision_rule(lbm_step.method, optimization=opt_params))
+ sim.write_files()
+
+ dir = '/localhome/seitz_local/projects/walberla/apps/autogen/'
+ os.makedirs(dir, exist_ok=True)
+ for k, v in ctx.files.items():
+ with open(join(dir, k), 'w') as file:
+ file.write(v)
+
+
+def test_resolve_parameters():
+ sym = TypedSymbol('s', create_type('double'))
+ sym2 = TypedSymbol('t', create_type('double'))
+
+ block_forest = UniformBlockforestFromConfig()
+
+ module = WalberlaModule(WalberlaMain(Block([
+ block_forest,
+ ResolveUndefinedSymbols(
+ Block([
+ SympyAssignment(sym, 1 + sym2),
+ ]), 'parameters')
+ ])))
+
+ print(module)
+