Bump waLBerla version requirements

CMakeLists.txt

@@ -603,7 +603,7 @@ if(ESPRESSO_BUILD_WITH_WALBERLA)
   FetchContent_Declare(
     walberla
     GIT_REPOSITORY https://i10git.cs.fau.de/walberla/walberla.git
-    GIT_TAG        0c8b4b926c6979288fd8a6846d02ec0870e1fe41
+    GIT_TAG        b0842e1a493ce19ef1bbb8d2cf382fc343970a7f
   )
   # workaround for https://gitlab.kitware.com/cmake/cmake/-/issues/21146
   if(NOT DEFINED walberla_SOURCE_DIR OR NOT EXISTS "${walberla_SOURCE_DIR}")

maintainer/walberla_kernels/Readme.md

@@ -20,7 +20,7 @@ The kernels can be regenerated with this shell script:
 
 ```sh
 # adapt these paths to the build environment
-export VERSION=1.2
+export VERSION=1.3.3
 export DEPS="${HOME}/walberla_deps"
 export PYTHONPATH="${DEPS}/${VERSION}/lbmpy:${DEPS}/${VERSION}/pystencils:${DEPS}/devel/walberla/python/"
 
@@ -49,6 +49,7 @@ generate_lb_kernels --single-precision
 generate_lb_kernels --gpu
 generate_lb_kernels --gpu --single-precision
 format_lb_kernels
+git diff src/walberla_bridge/src/lattice_boltzmann/generated_kernels/Dynamic_UBB_*CUDA*.cu # verify pragmas
 
 # EK kernels
 cd $(git rev-parse --show-toplevel)/src/walberla_bridge/src/electrokinetics/generated_kernels/

maintainer/walberla_kernels/code_generation_context.py

@@ -19,55 +19,11 @@
 
 import os
 import re
-import jinja2
 import hashlib
 import lbmpy
 import lbmpy_walberla
 import pystencils
 import pystencils_walberla
-import pystencils_walberla.codegen
-
-
-def adapt_pystencils():
-    """
-    Adapt pystencils to the SFINAE method (add the block offset lambda
-    callback and the time_step increment).
-    """
-    old_add_pystencils_filters_to_jinja_env = pystencils_walberla.codegen.add_pystencils_filters_to_jinja_env
-
-    def new_add_pystencils_filters_to_jinja_env(jinja_env):
-        # save original pystencils to adapt
-        old_add_pystencils_filters_to_jinja_env(jinja_env)
-        old_generate_members = jinja_env.filters["generate_members"]
-        old_generate_refs_for_kernel_parameters = jinja_env.filters[
-            "generate_refs_for_kernel_parameters"]
-
-        @jinja2.pass_context
-        def new_generate_members(*args, **kwargs):
-            output = old_generate_members(*args, **kwargs)
-            token = " block_offset_0_;"
-            if token in output:
-                i = output.index(token)
-                vartype = output[:i].split("\n")[-1].strip()
-                output += f"\nstd::function<void(IBlock *, {vartype}&, {vartype}&, {vartype}&)> block_offset_generator = [](IBlock * const, {vartype}&, {vartype}&, {vartype}&) {{ }};"  # nopep8
-            return output
-
-        def new_generate_refs_for_kernel_parameters(*args, **kwargs):
-            output = old_generate_refs_for_kernel_parameters(*args, **kwargs)
-            if "block_offset_0" in output:
-                old_token = "auto & block_offset_"
-                new_token = "auto block_offset_"
-                assert output.count(old_token) == 3, \
-                    f"could not find '{old_token}' in '''\n{output}\n'''"
-                output = output.replace(old_token, new_token)
-                output += "\nblock_offset_generator(block, block_offset_0, block_offset_1, block_offset_2);"
-            return output
-
-        # replace pystencils
-        jinja_env.filters["generate_members"] = new_generate_members
-        jinja_env.filters["generate_refs_for_kernel_parameters"] = new_generate_refs_for_kernel_parameters
-
-    pystencils_walberla.codegen.add_pystencils_filters_to_jinja_env = new_add_pystencils_filters_to_jinja_env
 
 
 def earmark_generated_kernels():
@@ -140,7 +96,6 @@ def __init__(self):
         sys.argv = sys.argv[:1]
         super().__init__()
         sys.argv = old_sys_argv
-        adapt_pystencils()
 
     def __exit__(self, *args, **kwargs):
         super().__exit__(*args, **kwargs)

maintainer/walberla_kernels/generate_ek_kernels.py

@@ -22,6 +22,7 @@
 import sympy as sp
 import lbmpy
 import argparse
+import packaging.specifiers
 
 import pystencils_espresso
 import code_generation_context
@@ -35,6 +36,12 @@
                     help='Use single-precision')
 args = parser.parse_args()
 
+# Make sure we have the correct versions of the required dependencies
+for module, requirement in [(ps, "==1.2"), (lbmpy, "==1.2")]:
+    assert packaging.specifiers.SpecifierSet(requirement).contains(module.__version__), \
+        f"{module.__name__} version {module.__version__} " \
+        f"doesn't match requirement {requirement}"
+
 double_precision: bool = not args.single_precision
 
 data_type_cpp = "double" if double_precision else "float"

maintainer/walberla_kernels/generate_lb_kernels.py

@@ -21,6 +21,7 @@
 import packaging.specifiers
 
 import sympy as sp
+import numpy as np
 
 import pystencils as ps
 import pystencils_walberla
@@ -52,7 +53,7 @@
     target = ps.Target.CPU
 
 # Make sure we have the correct versions of the required dependencies
-for module, requirement in [(ps, "==1.2"), (lbmpy, "==1.2")]:
+for module, requirement in [(ps, "==1.3.3"), (lbmpy, "==1.3.3")]:
     assert packaging.specifiers.SpecifierSet(requirement).contains(module.__version__), \
         f"{module.__name__} version {module.__version__} " \
         f"doesn't match requirement {requirement}"
@@ -117,7 +118,7 @@ def paramlist(parameters, keys):
 
     # generate initial densities
     for params, target_suffix in paramlist(parameters, (default_key,)):
-        pystencils_walberla.codegen.generate_sweep(
+        pystencils_walberla.generate_sweep(
             ctx,
             f"InitialPDFsSetter{precision_prefix}{target_suffix}",
             pystencils_espresso.generate_setters(ctx, method, params),
@@ -148,38 +149,44 @@ def paramlist(parameters, keys):
             params
         )
 
+    block_offsets = tuple(
+        ps.TypedSymbol(f"block_offset_{i}", np.uint32)
+        for i in range(3))
+
     # generate thermalized LB
     collision_rule_thermalized = lbmpy.creationfunctions.create_lb_collision_rule(
         method,
         zero_centered=False,
         fluctuating={
             "temperature": kT,
-            "block_offsets": "walberla",
+            "block_offsets": block_offsets,
             "rng_node": precision_rng
         },
         optimization={"cse_global": True,
                       "double_precision": ctx.double_accuracy}
     )
     for params, target_suffix in paramlist(parameters, ("GPU", "CPU", "AVX")):
+        stem = f"CollideSweep{precision_prefix}Thermalized{target_suffix}"
         pystencils_espresso.generate_collision_sweep(
             ctx,
             method,
             collision_rule_thermalized,
-            f"CollideSweep{precision_prefix}Thermalized{target_suffix}",
-            params
+            stem,
+            params,
+            block_offset=block_offsets,
         )
 
     # generate accessors
     for _, target_suffix in paramlist(parameters, ("GPU", "CPU")):
-        filename = f"FieldAccessors{precision_prefix}{target_suffix}"
+        stem = f"FieldAccessors{precision_prefix}{target_suffix}"
         if target == ps.Target.GPU:
             templates = {
-                f"{filename}.cuh": "templates/FieldAccessors.tmpl.cuh",
-                f"{filename}.cu": "templates/FieldAccessors.tmpl.cu",
+                f"{stem}.cuh": "templates/FieldAccessors.tmpl.cuh",
+                f"{stem}.cu": "templates/FieldAccessors.tmpl.cu",
             }
         else:
             templates = {
-                f"{filename}.h": "templates/FieldAccessors.tmpl.h",
+                f"{stem}.h": "templates/FieldAccessors.tmpl.h",
             }
         walberla_lbm_generation.generate_macroscopic_values_accessors(
             ctx, config, method, templates

maintainer/walberla_kernels/lbmpy_espresso.py

@@ -19,8 +19,8 @@
 
 import pystencils as ps
 
-import lbmpy.advanced_streaming.indexing
 import lbmpy.boundaries
+import lbmpy.custom_code_nodes
 
 import lbmpy_walberla.additional_data_handler
 
@@ -39,15 +39,20 @@ def data_initialisation(self, direction):
         This way, the dynamic UBB can be used to implement a LB boundary.
         '''
         code = super().data_initialisation(direction)
-        dirVec = self.stencil_info[direction][1]
-        token = ' = elementInitaliser(Cell(it.x(){}, it.y(){}, it.z(){}),'
-        old_initialiser = token.format('', '', '')
-        assert old_initialiser in code
-        new_initialiser = token.format(
-            '+' + str(dirVec[0]),
-            '+' + str(dirVec[1]),
-            '+' + str(dirVec[2])).replace('+-', '-')
-        return code.replace(old_initialiser, new_initialiser)
+        assert "InitialisationAdditionalData" in code
+        assert "elementInitialiser" in code
+        assert "element.vel_0" in code
+        bb_vec = self.stencil_info[direction][1]
+        cell_args = [f"it.{direction}() + {bb_vec[i]}".replace('+ -', '-')
+                     for i, direction in enumerate("xyz")]
+        code = [
+            "auto const InitialisationAdditionalData = elementInitialiser(",
+            f"Cell({', '.join(cell_args)}), blocks, *block);",
+            "element.vel_0 = InitialisationAdditionalData[0];",
+            "element.vel_1 = InitialisationAdditionalData[1];",
+            "element.vel_2 = InitialisationAdditionalData[2];",
+        ]
+        return "\n".join(code)
 
 
 class UBB(lbmpy.boundaries.UBB):
@@ -71,7 +76,7 @@ def __call__(self, f_out, f_in, dir_symbol,
         if len(assignments) > 1:
             out.extend(assignments[:-1])
 
-        neighbor_offset = lbmpy.advanced_streaming.indexing.NeighbourOffsetArrays.neighbour_offset(
+        neighbor_offset = lbmpy.custom_code_nodes.NeighbourOffsetArrays.neighbour_offset(
             dir_symbol, lb_method.stencil)
 
         assignment = assignments[-1]

maintainer/walberla_kernels/pystencils_espresso.py

@@ -23,6 +23,7 @@
 import lbmpy.updatekernels
 import pystencils as ps
 import pystencils_walberla
+import pystencils_walberla.utility
 
 
 def skip_philox_unthermalized(code, result_symbols, rng_name):
@@ -107,11 +108,11 @@ def generate_fields(config, stencil):
 
 
 def generate_config(ctx, params):
-    return pystencils_walberla.codegen.config_from_context(ctx, **params)
+    return pystencils_walberla.utility.config_from_context(ctx, **params)
 
 
 def generate_collision_sweep(
-        ctx, lb_method, collision_rule, class_name, params):
+        ctx, lb_method, collision_rule, class_name, params, **kwargs):
     config = generate_config(ctx, params)
 
     # Symbols for PDF (twice, due to double buffering)
@@ -127,8 +128,8 @@ def generate_collision_sweep(
         collide_update_rule, config=config, **params)
     collide_ast.function_name = 'kernel_collide'
     collide_ast.assumed_inner_stride_one = True
-    pystencils_walberla.codegen.generate_sweep(
-        ctx, class_name, collide_ast, **params)
+    pystencils_walberla.generate_sweep(
+        ctx, class_name, collide_ast, **params, **kwargs)
 
 
 def generate_stream_sweep(ctx, lb_method, class_name, params):
@@ -144,7 +145,7 @@ def generate_stream_sweep(ctx, lb_method, class_name, params):
     stream_ast = ps.create_kernel(stream_update_rule, config=config, **params)
     stream_ast.function_name = 'kernel_stream'
     stream_ast.assumed_inner_stride_one = True
-    pystencils_walberla.codegen.generate_sweep(
+    pystencils_walberla.generate_sweep(
         ctx, class_name, stream_ast,
         field_swaps=[(fields['pdfs'], fields['pdfs_tmp'])], **params)