Allocating multiple blocks to one mpi rank in LB

doc/sphinx/lb.rst

@@ -214,7 +214,8 @@ will be used instead of the default ones.
 
 .. note::
 
-    At the moment, LB only supports the case ``shear_plane_normal="y"``.
+    At the moment, LB only supports the case ``shear_plane_normal="y"`` and
+    doesn't allow domain decomposition along the shear and normal directions.
 
 .. _Reading and setting properties of single lattice nodes:
 

maintainer/benchmarks/lb.py

@@ -50,6 +50,11 @@
 parser.add_argument("--output", metavar="FILEPATH", action="store",
                     type=str, required=False, default="benchmarks.csv",
                     help="Output file (default: benchmarks.csv)")
+parser.add_argument("--blocks_per_mpi_rank", action="store", nargs=3,
+                    type=int, default=[1, 1, 1], required=False,
+                    help="blocks per mpi rank")
+parser.add_argument("--weak_scaling", action="store_true", required=False,
+                    help="The measurement of weak scaling")
 
 args = parser.parse_args()
 
@@ -101,6 +106,9 @@
     lb_grid = 3 * [lb_grid]
     box_l = 3 * [box_l]
 
+if args.weak_scaling:
+    box_l *= system.cell_system.node_grid
+
 print(f"box length: {box_l}")
 print(f"LB shape: {lb_grid}")
 print(f"LB agrid: {agrid:.3f}")
@@ -145,7 +153,8 @@
 if args.multi_gpu:
     system.cuda_init_handle.call_method("set_device_id_per_rank")
 lbf = lb_class(agrid=agrid, tau=system.time_step, kinematic_viscosity=1.,
-               density=1., single_precision=args.single_precision)
+               density=1., single_precision=args.single_precision,
+               blocks_per_mpi_rank=args.blocks_per_mpi_rank)
 system.lb = lbf
 if n_part:
     system.thermostat.set_lb(LB_fluid=lbf, gamma=1., seed=42)

src/core/unit_tests/ek_interface_test.cpp

@@ -83,7 +83,8 @@ static auto make_ek_actor() {
   auto constexpr n_ghost_layers = 1u;
   auto constexpr single_precision = true;
   ek_lattice = std::make_shared<LatticeWalberla>(
-      params.grid_dimensions, ::communicator.node_grid, n_ghost_layers);
+      params.grid_dimensions, ::communicator.node_grid,
+      ::communicator.node_grid, n_ghost_layers);
   ek_container = std::make_shared<EK::EKWalberla::ek_container_type>(
       params.tau, walberla::new_ek_poisson_none(ek_lattice, single_precision));
   ek_reactions = std::make_shared<EK::EKWalberla::ek_reactions_type>();

src/core/unit_tests/lb_particle_coupling_test.cpp

@@ -102,7 +102,8 @@ static auto make_lb_actor() {
   auto constexpr single_precision = false;
   lb_params = std::make_shared<LB::LBWalberlaParams>(params.agrid, params.tau);
   lb_lattice = std::make_shared<LatticeWalberla>(
-      params.grid_dimensions, ::communicator.node_grid, n_ghost_layers);
+      params.grid_dimensions, ::communicator.node_grid,
+      ::communicator.node_grid, n_ghost_layers);
   lb_fluid = new_lb_walberla_cpu(lb_lattice, params.viscosity, params.density,
                                  single_precision);
   lb_fluid->set_collision_model(params.kT, params.seed);
@@ -535,7 +536,8 @@ bool test_lb_domain_mismatch_local() {
   auto const params = std::make_shared<LB::LBWalberlaParams>(0.5, 0.01);
   ::communicator.node_grid = node_grid_reversed;
   auto const lattice = std::make_shared<LatticeWalberla>(
-      Utils::Vector3i{12, 12, 12}, node_grid_original, n_ghost_layers);
+      Utils::Vector3i{12, 12, 12}, node_grid_original, node_grid_original,
+      n_ghost_layers);
   auto const ptr = new_lb_walberla_cpu(lattice, 1.0, 1.0, false);
   ptr->set_collision_model(0.0, 0);
   ::communicator.node_grid = node_grid_original;

src/python/espressomd/detail/walberla.py

@@ -22,22 +22,31 @@
 import numpy as np
 
 import espressomd.shapes
-import espressomd.code_features
 from espressomd.script_interface import ScriptInterfaceHelper, script_interface_register
 
 
 @script_interface_register
 class LatticeWalberla(ScriptInterfaceHelper):
     """
-    Interface to a waBLerla lattice.
+    Interface to a waLBerla lattice.
+
+    Parameters
+    ----------
+    agrid : :obj:`float`
+        Lattice constant. The box size in every direction must be an integer
+        multiple of ``agrid``. Cannot be provided together with ``lattice``.
+    n_ghost_layers : :obj:`int`, optional
+        Lattice ghost layer thickness in units of ``agrid``.
+    blocks_per_mpi_rank : (3,) array_like of :obj:`int`, optional
+        Distribute more than one block to each MPI rank.
+        Meant to improve cache locality. Experimental.
+
     """
     _so_name = "walberla::LatticeWalberla"
     _so_creation_policy = "GLOBAL"
+    _so_features = ("WALBERLA",)
 
     def __init__(self, *args, **kwargs):
-        if not espressomd.code_features.has_features("WALBERLA"):
-            raise NotImplementedError("Feature WALBERLA not compiled in")
-
         if "sip" not in kwargs:
             params = self.default_params()
             params.update(kwargs)
@@ -46,14 +55,17 @@ def __init__(self, *args, **kwargs):
         else:
             super().__init__(**kwargs)
 
-    def valid_keys(self):
-        return {"agrid", "n_ghost_layers"}
+    @classmethod
+    def valid_keys(cls):
+        return {"agrid", "n_ghost_layers", "blocks_per_mpi_rank"}
 
-    def required_keys(self):
-        return self.valid_keys()
+    @classmethod
+    def required_keys(cls):
+        return {"agrid"}
 
-    def default_params(self):
-        return {}
+    @classmethod
+    def default_params(cls):
+        return {"n_ghost_layers": 1, "blocks_per_mpi_rank": [1, 1, 1]}
 
     def get_node_indices_inside_shape(self, shape):
         if not isinstance(shape, espressomd.shapes.Shape):
@@ -146,10 +158,9 @@ def get_slice_bounding_box(slices, grid_size):
 
 
 class VTKOutputBase(ScriptInterfaceHelper):
+    _so_features = ("WALBERLA",)
 
     def __init__(self, *args, **kwargs):
-        if not espressomd.code_features.has_features("WALBERLA"):
-            raise NotImplementedError("Feature WALBERLA not compiled in")
         if "sip" not in kwargs:
             params = self.default_params()
             params.update(kwargs)

src/python/espressomd/electrokinetics.py

@@ -27,6 +27,12 @@
 import espressomd.shapes
 
 
+def _check_lattice_blocks(class_name, pack):
+    if "lattice" in pack and np.prod(pack["lattice"].blocks_per_mpi_rank) != 1:
+        raise NotImplementedError(
+            f"Using more than one block per MPI rank is not supported for {class_name}")
+
+
 @script_interface_register
 class EKFFT(ScriptInterfaceHelper):
     """
@@ -46,6 +52,10 @@ class EKFFT(ScriptInterfaceHelper):
     _so_features = ("WALBERLA_FFT",)
     _so_creation_policy = "GLOBAL"
 
+    def __init__(self, *args, **kwargs):
+        _check_lattice_blocks(self.__class__.__name__, kwargs)
+        super().__init__(*args, **kwargs)
+
 
 @script_interface_register
 class EKNone(ScriptInterfaceHelper):
@@ -64,6 +74,10 @@ class EKNone(ScriptInterfaceHelper):
     _so_features = ("WALBERLA",)
     _so_creation_policy = "GLOBAL"
 
+    def __init__(self, *args, **kwargs):
+        _check_lattice_blocks(self.__class__.__name__, kwargs)
+        super().__init__(*args, **kwargs)
+
 
 @script_interface_register
 class EKSpecies(ScriptInterfaceHelper,
@@ -167,6 +181,7 @@ def __init__(self, *args, **kwargs):
         if "sip" not in kwargs:
             params = self.default_params()
             params.update(kwargs)
+            _check_lattice_blocks(self.__class__.__name__, params)
             super().__init__(*args, **params)
         else:
             super().__init__(**kwargs)
@@ -580,6 +595,7 @@ class VTKOutput(VTKOutputBase):
     _so_name = "walberla::EKVTKHandle"
     _so_creation_policy = "GLOBAL"
     _so_bind_methods = ("enable", "disable", "write")
+    _so_features = ("WALBERLA",)
 
     def required_keys(self):
         return self.valid_keys() - self.default_params().keys()

src/python/espressomd/lb.py

@@ -57,8 +57,8 @@ def validate_params(self, params):
         pass
 
     def valid_keys(self):
-        return {"agrid", "tau", "density", "ext_force_density",
-                "kinematic_viscosity", "lattice", "kT", "seed"}
+        return {"agrid", "tau", "lattice", "density", "ext_force_density",
+                "kinematic_viscosity", "kT", "seed", "blocks_per_mpi_rank"}
 
     def required_keys(self):
         return {"lattice", "density", "kinematic_viscosity", "tau"}
@@ -119,7 +119,7 @@ class LBFluidWalberla(HydrodynamicInteraction,
 
     Parameters
     ----------
-    lattice : :obj:`espressomd.lb.LatticeWalberla <espressomd.detail.walberla.LatticeWalberla>`
+    lattice : :obj:`~espressomd.detail.walberla.LatticeWalberla`
         Lattice object. If not provided, a default one will be constructed
         using the ``agrid`` parameter.
     agrid : :obj:`float`
@@ -141,6 +141,9 @@ class LBFluidWalberla(HydrodynamicInteraction,
         Required for a thermalized fluid. Must be positive.
     single_precision : :obj:`bool`, optional
         Use single-precision floating-point arithmetic.
+    \\*\\*kwargs :
+        Additional parameters forwarded to the
+        :obj:`~espressomd.detail.walberla.LatticeWalberla` constructor.
 
     Methods
     -------
@@ -235,14 +238,21 @@ def __init__(self, *args, **kwargs):
     def validate_params(self, params):
         super().validate_params(params)
 
+        # extract lattice-specific parameters
+        lattice_params = {}
+        for key in LatticeWalberla.valid_keys():
+            if key in params:
+                lattice_params[key] = params.pop(key)
+
         # construct default lattice if necessary
         if params.get("lattice") is None:
-            if "agrid" not in params:
-                raise ValueError("missing argument 'lattice' or 'agrid'")
-            params["lattice"] = LatticeWalberla(
-                agrid=params.pop("agrid"), n_ghost_layers=1)
-        elif "agrid" in params:
-            raise ValueError("cannot provide both 'lattice' and 'agrid'")
+            for key in LatticeWalberla.required_keys():
+                if key not in lattice_params:
+                    raise ValueError(f"missing argument 'lattice' or '{key}'")
+            params["lattice"] = LatticeWalberla(**lattice_params)
+        elif lattice_params:
+            any_key = list(lattice_params.keys())[0]
+            raise ValueError(f"cannot provide both 'lattice' and '{any_key}'")
 
         utils.check_required_keys(self.required_keys(), params.keys())
         utils.check_valid_keys(self.valid_keys(), params.keys())
@@ -660,6 +670,7 @@ class VTKOutput(VTKOutputBase):
     _so_name = "walberla::LBVTKHandle"
     _so_creation_policy = "GLOBAL"
     _so_bind_methods = ("enable", "disable", "write")
+    _so_features = ("WALBERLA",)
 
     def required_keys(self):
         return self.valid_keys() - self.default_params().keys()

src/script_interface/walberla/LBFluid.cpp

@@ -139,6 +139,12 @@ void LBFluidGPU::make_instance(VariantMap const &params) {
   auto const visc = get_value<double>(params, "kinematic_viscosity");
   auto const dens = get_value<double>(params, "density");
   auto const precision = get_value<bool>(params, "single_precision");
+  auto const blocks_per_mpi_rank = get_value<Utils::Vector3i>(
+      m_lattice->get_parameter("blocks_per_mpi_rank"));
+  if (blocks_per_mpi_rank != Utils::Vector3i{{1, 1, 1}}) {
+    throw std::runtime_error(
+        "Using more than one block per MPI rank is not supported for GPU LB");
+  }
   auto const lb_lattice = m_lattice->lattice();
   auto const lb_visc = m_conv_visc * visc;
   auto const lb_dens = m_conv_dens * dens;

src/script_interface/walberla/LatticeWalberla.hpp

@@ -43,6 +43,7 @@ class LatticeWalberla : public AutoParameters<LatticeWalberla> {
   std::shared_ptr<::LatticeWalberla> m_lattice;
   double m_agrid;
   Utils::Vector3d m_box_l;
+  Utils::Vector3i m_blocks_per_mpi_rank;
 
 public:
   LatticeWalberla() {
@@ -53,14 +54,20 @@ class LatticeWalberla : public AutoParameters<LatticeWalberla> {
         {"shape", AutoParameter::read_only,
          [this]() { return m_lattice->get_grid_dimensions(); }},
         {"_box_l", AutoParameter::read_only, [this]() { return m_box_l; }},
+        {"blocks_per_mpi_rank", AutoParameter::read_only,
+         [this]() { return m_blocks_per_mpi_rank; }},
     });
   }
 
   void do_construct(VariantMap const &args) override {
     auto const &box_geo = *::System::get_system().box_geo;
     m_agrid = get_value<double>(args, "agrid");
     m_box_l = get_value_or<Utils::Vector3d>(args, "_box_l", box_geo.length());
+    m_blocks_per_mpi_rank =
+        get_value<Utils::Vector3i>(args, "blocks_per_mpi_rank");
     auto const n_ghost_layers = get_value<int>(args, "n_ghost_layers");
+    auto const block_grid = Utils::hadamard_product(::communicator.node_grid,
+                                                    m_blocks_per_mpi_rank);
 
     context()->parallel_try_catch([&]() {
       if (m_agrid <= 0.) {
@@ -69,10 +76,13 @@ class LatticeWalberla : public AutoParameters<LatticeWalberla> {
       if (n_ghost_layers < 0) {
         throw std::domain_error("Parameter 'n_ghost_layers' must be >= 0");
       }
+      if (not(m_blocks_per_mpi_rank >= Utils::Vector3i::broadcast(1))) {
+        throw std::domain_error("Parameter 'blocks_per_mpi_rank' must be >= 1");
+      }
       auto const grid_dim =
           ::LatticeWalberla::calc_grid_dimensions(m_box_l, m_agrid);
       m_lattice = std::make_shared<::LatticeWalberla>(
-          grid_dim, ::communicator.node_grid,
+          grid_dim, ::communicator.node_grid, block_grid,
           static_cast<unsigned int>(n_ghost_layers));
     });
   }

src/walberla_bridge/CMakeLists.txt

@@ -37,7 +37,6 @@ list(APPEND SKIP_CLANG_TIDY_CHECKS "-modernize-use-equals-delete")
 list(APPEND SKIP_CLANG_TIDY_CHECKS "-modernize-return-braced-init-list")
 list(APPEND SKIP_CLANG_TIDY_CHECKS "-bugprone-crtp-constructor-accessibility")
 list(APPEND SKIP_CLANG_TIDY_CHECKS "-bugprone-narrowing-conversions")
-list(APPEND SKIP_CLANG_TIDY_CHECKS "-bugprone-assert-side-effect")
 list(APPEND SKIP_CLANG_TIDY_CHECKS "-bugprone-exception-escape")
 list(APPEND SKIP_CLANG_TIDY_CHECKS "-bugprone-branch-clone")
 if(WALBERLA_BUILD_WITH_CUDA)