PackInfo code generation (#4994)

Fixes #4921

Description of changes:
- reduce volume of the LB communication during the streaming step
- avoid superfluous LB ghost communication outside the streaming step
- use AVX streaming kernels

Author: kodiakhq[bot]

maintainer/benchmarks/CMakeLists.txt

@@ -150,4 +150,5 @@ add_custom_target(
   COMMAND ${CMAKE_CTEST_COMMAND} --timeout ${ESPRESSO_TEST_TIMEOUT}
           ${ESPRESSO_CTEST_ARGS} --output-on-failure)
 
-add_dependencies(benchmark benchmark_python benchmarks_data)
+add_dependencies(benchmark_python pypresso benchmarks_data)
+add_dependencies(benchmark benchmark_python)

maintainer/benchmarks/benchmarks.py

@@ -84,7 +84,7 @@ def get_timings(system, n_steps, n_iterations, verbose=True):
             energy = system.analysis.energy()["total"]
             verlet = system.cell_system.get_state()["verlet_reuse"]
             print(
-                f"step {i}, time: {1000 * t:.1f} ms, verlet: {verlet:.2f}, energy: {energy:.2e}")
+                f"step {i}, time: {1000 * t:.2f} ms, verlet: {verlet:.2f}, energy: {energy:.2e}")
     return np.array(timings)
 
 

maintainer/benchmarks/lb.py

@@ -149,7 +149,7 @@
 
 # average time
 avg, ci = benchmarks.get_average_time(timings)
-print(f"average: {1000 * avg:.1f} +/- {1000 * ci:.1f} ms (95% C.I.)")
+print(f"average: {1000 * avg:.2f} +/- {1000 * ci:.2f} ms (95% C.I.)")
 
 # write report
 benchmarks.write_report(args.output, n_proc, timings, measurement_steps)

maintainer/walberla_kernels/generate_lb_kernels.py

@@ -98,6 +98,8 @@ def paramlist(parameters, keys):
     stencil = lbmpy.stencils.LBStencil(lbmpy.enums.Stencil.D3Q19)
     fields = pystencils_espresso.generate_fields(config, stencil)
     force_field = fields["force"]
+    lbm_opt = lbmpy.LBMOptimisation(symbolic_field=fields["pdfs"])
+    streaming_pattern = "push"
 
     # LB Method definition
     method = lbmpy.creationfunctions.create_mrt_orthogonal(
@@ -133,12 +135,11 @@ def paramlist(parameters, keys):
                                 force_model=lbmpy.ForceModel.GUO,
                                 force=force_field.center_vector,
                                 kernel_type="collide_only")
-    lbm_opt = lbmpy.LBMOptimisation(symbolic_field=fields["pdfs"])
-    le_collision_rule_unthermalized = lbmpy.create_lb_update_rule(
+    le_update_rule_unthermalized = lbmpy.create_lb_update_rule(
         lbm_config=le_config,
         lbm_optimisation=lbm_opt)
     le_collision_rule_unthermalized = lees_edwards.add_lees_edwards_to_collision(
-        config, le_collision_rule_unthermalized,
+        config, le_update_rule_unthermalized,
         fields["pdfs"], stencil, 1)  # shear_dir_normal y
     for params, target_suffix in paramlist(parameters, ("GPU", "CPU", "AVX")):
         pystencils_espresso.generate_collision_sweep(
@@ -153,8 +154,8 @@ def paramlist(parameters, keys):
         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(
+    # generate thermalized LB collision rule
+    lb_collision_rule_thermalized = lbmpy.creationfunctions.create_lb_collision_rule(
         method,
         zero_centered=False,
         fluctuating={
@@ -170,7 +171,7 @@ def paramlist(parameters, keys):
         pystencils_espresso.generate_collision_sweep(
             ctx,
             method,
-            collision_rule_thermalized,
+            lb_collision_rule_thermalized,
             stem,
             params,
             block_offset=block_offsets,
@@ -192,6 +193,30 @@ def paramlist(parameters, keys):
             ctx, config, method, templates
         )
 
+    # generate PackInfo
+    assignments = pystencils_espresso.generate_pack_info_pdfs_field_assignments(
+        fields, streaming_pattern="pull")
+    spec = pystencils_espresso.generate_pack_info_vector_field_specifications(
+        config, stencil, force_field.layout)
+    for params, target_suffix in paramlist(parameters, ["CPU"]):
+        pystencils_walberla.generate_pack_info_from_kernel(
+            ctx, f"PackInfoPdf{precision_prefix}{target_suffix}", assignments,
+            kind="pull", **params)
+        pystencils_walberla.generate_pack_info(
+            ctx, f"PackInfoVec{precision_prefix}{target_suffix}", spec, **params)
+        if target_suffix == "CUDA":
+            continue
+        token = "\n       //TODO: optimize by generating kernel for this case\n"
+        for field_suffix in ["Pdf", "Vec"]:
+            class_name = f"PackInfo{field_suffix}{precision_prefix}{target_suffix}"  # nopep8
+            with open(f"{class_name}.h", "r+") as f:
+                content = f.read()
+                assert token in content
+                content = content.replace(token, "\n")
+                f.seek(0)
+                f.truncate()
+                f.write(content)
+
     # boundary conditions
     ubb_dynamic = lbmpy_espresso.UBB(
         lambda *args: None, dim=3, data_type=config.data_type.default_factory())
@@ -202,7 +227,7 @@ def paramlist(parameters, keys):
         lbmpy_walberla.generate_boundary(
             ctx, f"Dynamic_UBB_{precision_suffix}{target_suffix}", ubb_dynamic,
             method, additional_data_handler=ubb_data_handler,
-            streaming_pattern="push", target=target)
+            streaming_pattern=streaming_pattern, target=target)
 
         with open(f"Dynamic_UBB_{precision_suffix}{target_suffix}.h", "r+") as f:
             content = f.read()

maintainer/walberla_kernels/pystencils_espresso.py

@@ -164,9 +164,8 @@ def __init__(self, dim, time_step=ps.typing.TypedSymbol(
 data_type_np = {'double': 'float64', 'float': 'float32'}
 
 
-def generate_fields(config, stencil):
+def generate_fields(config, stencil, field_layout='fzyx'):
     dtype = data_type_np[config.data_type.default_factory().c_name]
-    field_layout = 'fzyx'
     q = len(stencil)
     dim = len(stencil[0])
 
@@ -208,6 +207,60 @@ def generate_fields(config, stencil):
     return fields
 
 
+def generate_pack_info_pdfs_field_assignments(fields, streaming_pattern):
+    """
+    Visualize the stencil directions with::
+
+       import lbmpy
+       import matplotlib.pyplot as plt
+       stencil = lbmpy.LBStencil(lbmpy.Stencil.D3Q19)
+       stencil.plot(data=[i for i in range(19)])
+       plt.show()
+
+    """
+    stencil = lbmpy.enums.Stencil.D3Q19
+    lbm_config = lbmpy.LBMConfig(stencil=stencil,
+                                 method=lbmpy.Method.CUMULANT,
+                                 compressible=True,
+                                 zero_centered=False,
+                                 weighted=True,
+                                 streaming_pattern=streaming_pattern,
+                                 relaxation_rate=sp.Symbol("omega_shear"),
+                                 )
+    lbm_opt = lbmpy.LBMOptimisation(
+        symbolic_field=fields["pdfs" if streaming_pattern ==
+                              "pull" else "pdfs_tmp"],
+        symbolic_temporary_field=fields["pdfs" if streaming_pattern ==
+                                        "push" else "pdfs_tmp"],
+        field_layout=fields['pdfs'].layout)
+    lbm_update_rule = lbmpy.create_lb_update_rule(
+        lbm_config=lbm_config,
+        lbm_optimisation=lbm_opt)
+    return lbm_update_rule.all_assignments
+
+
+def generate_pack_info_vector_field_specifications(config, stencil, layout):
+    import collections
+    import itertools
+    field = ps.Field.create_generic(
+        "field",
+        3,
+        data_type_np[config.data_type.default_factory().c_name],
+        index_dimensions=1,
+        layout=layout,
+        index_shape=(3,)
+    )
+    q = len(stencil)
+    coord = itertools.product(*[(-1, 0, 1)] * 3)
+    if q == 19:
+        dirs = tuple((i, j, k) for i, j, k in coord if i**2 + j**2 + k**2 != 3)
+    else:
+        dirs = tuple((i, j, k) for i, j, k in coord)
+    spec = collections.defaultdict(set)
+    spec[dirs] = {field[0, 0, 0](i) for i in range(3)}
+    return spec
+
+
 def generate_config(ctx, params):
     return pystencils_walberla.utility.config_from_context(ctx, **params)
 

src/core/integrate.cpp

@@ -629,6 +629,7 @@ int System::System::integrate(int n_steps, int reuse_forces) {
           propagation.lb_skipped_md_steps = 0;
           propagation.ek_skipped_md_steps = 0;
           lb.propagate();
+          lb.ghost_communication_vel();
           ek.propagate();
         }
       } else if (lb_active) {
@@ -654,6 +655,9 @@ int System::System::integrate(int n_steps, int reuse_forces) {
 #ifdef VIRTUAL_SITES_INERTIALESS_TRACERS
       if (thermostat->lb and
           (propagation.used_propagations & PropagationMode::TRANS_LB_TRACER)) {
+        if (lb_active) {
+          lb.ghost_communication_vel();
+        }
         lb_tracers_propagate(*cell_structure, lb, time_step);
       }
 #endif
@@ -678,6 +682,9 @@ int System::System::integrate(int n_steps, int reuse_forces) {
     }
 
   } // for-loop over integration steps
+  if (lb_active) {
+    lb.ghost_communication();
+  }
   lees_edwards->update_box_params(*box_geo, sim_time);
 #ifdef CALIPER
   CALI_CXX_MARK_LOOP_END(integration_loop);

src/core/lb/LBNone.hpp

@@ -29,6 +29,9 @@ namespace LB {
 
 struct LBNone {
   void propagate() { throw NoLBActive{}; }
+  void ghost_communication() { throw NoLBActive{}; }
+  void ghost_communication_pdf() { throw NoLBActive{}; }
+  void ghost_communication_vel() { throw NoLBActive{}; }
   double get_agrid() const { throw NoLBActive{}; }
   double get_tau() const { throw NoLBActive{}; }
   double get_kT() const { throw NoLBActive{}; }

src/core/lb/LBWalberla.cpp

@@ -52,6 +52,16 @@ Utils::VectorXd<9> LBWalberla::get_pressure_tensor() const {
 
 void LBWalberla::propagate() { lb_fluid->integrate(); }
 
+void LBWalberla::ghost_communication() { lb_fluid->ghost_communication(); }
+
+void LBWalberla::ghost_communication_pdf() {
+  lb_fluid->ghost_communication_vel();
+}
+
+void LBWalberla::ghost_communication_vel() {
+  lb_fluid->ghost_communication_vel();
+}
+
 void LBWalberla::lebc_sanity_checks(unsigned int shear_direction,
                                     unsigned int shear_plane_normal) const {
   lb_fluid->check_lebc(shear_direction, shear_plane_normal);

src/core/lb/LBWalberla.hpp

@@ -72,6 +72,9 @@ struct LBWalberla {
   std::vector<Utils::Vector3d>
   get_velocities_at_pos(std::vector<Utils::Vector3d> const &pos);
   void propagate();
+  void ghost_communication();
+  void ghost_communication_pdf();
+  void ghost_communication_vel();
   void veto_time_step(double time_step) const;
   void veto_kT(double kT) const;
   void sanity_checks(System::System const &system) const;

src/core/lb/Solver.cpp

@@ -73,6 +73,21 @@ void Solver::propagate() {
   std::visit([](auto &ptr) { ptr->propagate(); }, *impl->solver);
 }
 
+void Solver::ghost_communication() {
+  check_solver(impl);
+  std::visit([](auto &ptr) { ptr->ghost_communication(); }, *impl->solver);
+}
+
+void Solver::ghost_communication_pdf() {
+  check_solver(impl);
+  std::visit([](auto &ptr) { ptr->ghost_communication_pdf(); }, *impl->solver);
+}
+
+void Solver::ghost_communication_vel() {
+  check_solver(impl);
+  std::visit([](auto &ptr) { ptr->ghost_communication_vel(); }, *impl->solver);
+}
+
 void Solver::sanity_checks() const {
   if (impl->solver) {
     auto const &system = get_system();