Update generated lattice model

python/lbmpy_walberla/templates/LatticeModel.tmpl.cpp

@@ -90,10 +90,22 @@ void {{class_name}}::Sweep::stream( IBlock * block, const uint_t numberOfGhostLa
     {{stream_kernel|generate_block_data_to_field_extraction(parameters=['pdfs', 'pdfs_tmp'])|indent(4)}}
 
     {{stream_kernel|generate_call('cell_idx_c(numberOfGhostLayersToInclude)')|indent(4)}}
-    
+
     {{stream_kernel|generate_swaps|indent(4)}}
 }
 
+// IMPORTANT REMARK:
+// This is specifically implemented for using generated kernels in the waLBerla's free surface LBM and is
+// implemented in rather unflexible fashion. Therefore, it should not be extended and in the long-term, the free
+// surface implementation should be refactored such that the general generated stream() is applicable.
+void {{class_name}}::Sweep::streamInCellInterval( {{stream_kernel|generate_field_type()}} * const pdfs,
+                                                  {{stream_kernel|generate_field_type()}} * pdfs_tmp,
+                                                  const CellInterval & ci )
+{
+    {{stream_kernel|generate_call(ghost_layers_to_include=0, cell_interval="ci", stream="stream")|indent(4)}}
+}
+
+
 
 } // namespace {{namespace}}
 } // namespace walberla

python/lbmpy_walberla/templates/LatticeModel.tmpl.h

@@ -118,6 +118,13 @@ public:
             streamCollide( block, numberOfGhostLayersToInclude );
         }
 
+        // IMPORTANT REMARK:
+        // This is specifically implemented for using generated kernels in the waLBerla's free surface LBM and is
+        // implemented in rather unflexible fashion. Therefore, it should not be extended and in the long-term, the free
+        // surface implementation should be refactored such that the general generated stream() is applicable.
+        void streamInCellInterval( {{stream_kernel|generate_field_type()}} * const pdfSrcField,
+                                   {{stream_kernel|generate_field_type()}} * pdfDstField, const CellInterval & ci );
+
     private:
         {{stream_collide_kernel|generate_members(only_fields=True)|indent(8)}}
     };

python/pystencils_walberla/jinja_filters.py

@@ -14,6 +14,8 @@ from pystencils.data_types import TypedSymbol, get_base_type
 from pystencils.field import FieldType
 from pystencils.sympyextensions import prod
 
+temporary_fieldPointerTemplate = """{type}"""
+
 temporary_fieldMemberTemplate = """
 private: std::set< {type} *, field::SwapableCompare< {type} * > > cache_{original_field_name}_;"""
 
@@ -269,8 +271,15 @@ def generate_call(ctx, kernel, ghost_layers_to_include=0, cell_interval=None, st
             return [-ghost_layers_to_include - required_ghost_layers] * field_object.spatial_dimensions
         else:
             assert ghost_layers_to_include == 0
-            return [sp.Symbol("{ci}.{coord}Min()".format(coord=coord_name, ci=cell_interval)) - required_ghost_layers
-                    for coord_name in ('x', 'y', 'z')]
+            if field_object.spatial_dimensions == 3:
+                return [sp.Symbol("{ci}.{coord}Min()".format(coord=coord_name, ci=cell_interval)) - required_ghost_layers
+                        for coord_name in ('x', 'y', 'z')]
+            elif field_object.spatial_dimensions == 2:
+                return [sp.Symbol("{ci}.{coord}Min()".format(coord=coord_name, ci=cell_interval)) - required_ghost_layers
+                        for coord_name in ('x', 'y')]
+            else:
+                raise NotImplementedError(f"Only 2D and 3D fields are supported but a field with "
+                                          f"{field_object.spatial_dimensions} dimensions was passed")
 
     def get_end_coordinates(field_object):
         if cell_interval is None:
@@ -447,6 +456,27 @@ def generate_destructor(kernel_info, class_name):
         return temporary_constructor.format(contents=contents, class_name=class_name)
 
 
+# IMPORTANT REMARK:
+# This is specifically implemented for using generated kernels in the waLBerla's free surface LBM and is
+# implemented in rather unflexible fashion. Therefore, it should not be extended and in the long-term, the free
+# surface implementation should be refactored such that the general generated stream() is applicable.
+@jinja2_context_decorator
+def generate_field_type(ctx, kernel_info):
+    fields = {f.name: f for f in kernel_info.fields_accessed}
+    target = translate_target(ctx['target'])
+    is_gpu = target == Target.GPU
+
+    result = []
+    for field_name in kernel_info.temporary_fields:
+        f = fields[field_name]
+        assert field_name.endswith('_tmp')
+        original_field_name = field_name[:-len('_tmp')]
+        f_size = get_field_fsize(f)
+        field_type = make_field_type(get_base_type(f.dtype), f_size, is_gpu)
+        result.append(temporary_fieldPointerTemplate.format(type=field_type, original_field_name=original_field_name))
+    return "\n".join(result)
+
+
 @jinja2_context_decorator
 def nested_class_method_definition_prefix(ctx, nested_class_name):
     outer_class = ctx['class_name']
@@ -517,6 +547,7 @@ def add_pystencils_filters_to_jinja_env(jinja_env):
     jinja_env.filters['generate_swaps'] = generate_swaps
     jinja_env.filters['generate_refs_for_kernel_parameters'] = generate_refs_for_kernel_parameters
     jinja_env.filters['generate_destructor'] = generate_destructor
+    jinja_env.filters['generate_field_type'] = generate_field_type
     jinja_env.filters['nested_class_method_definition_prefix'] = nested_class_method_definition_prefix
     jinja_env.filters['type_identifier_list'] = type_identifier_list
     jinja_env.filters['identifier_list'] = identifier_list

tests/lbm/CMakeLists.txt

@@ -156,4 +156,14 @@ waLBerla_compile_test( FILES codegen/InplaceStreamingCodegenTest.cpp
 waLBerla_execute_test( NAME InplaceStreamingCodegenTest
                        COMMAND $<TARGET_FILE:InplaceStreamingCodegenTest> ${CMAKE_CURRENT_SOURCE_DIR}/codegen/InplaceStreamingCodegen3D.prm  )
 
+
+waLBerla_generate_target_from_python( NAME StreamInCellIntervalCodegen
+                                      FILE codegen/StreamInCellIntervalCodegen.py
+                                      OUT_FILES GeneratedLatticeModel.h GeneratedLatticeModel.cpp)
+
+waLBerla_compile_test( FILES    codegen/StreamInCellIntervalCodegenTest.cpp
+                       DEPENDS  blockforest field lbm timeloop StreamInCellIntervalCodegen )
+
+waLBerla_execute_test( NAME     StreamInCellIntervalCodegenTest )
+
 endif()

tests/lbm/codegen/StreamInCellIntervalCodegen.py

+import sympy as sp
+
+from lbmpy.creationfunctions import create_lb_collision_rule, create_lb_update_rule
+from lbmpy import LBMConfig
+
+from pystencils_walberla import CodeGeneration
+from lbmpy_walberla import generate_lattice_model
+
+# general parameters
+stencil = 'D3Q19'
+omega = sp.Symbol('omega')
+layout = 'fzyx'
+
+# optimizations to be used by the code generator
+optimizations = {'cse_global': True, 'field_layout': layout}
+
+# method definition
+method_params = {'stencil': stencil,
+                 'method': 'srt',
+                 'relaxation_rate': omega,
+                 'compressible': True}
+
+lbm_config = LBMConfig(streaming_pattern='pull')
+
+collision_rule = create_lb_collision_rule(optimization=optimizations, **method_params)
+update_rule = create_lb_update_rule(collision_rule=collision_rule, lbm_config=lbm_config, optimization=optimizations)
+
+with CodeGeneration() as ctx:
+    generate_lattice_model(ctx, "GeneratedLatticeModel", collision_rule, field_layout=layout)

tests/lbm/codegen/StreamInCellIntervalCodegenTest.cpp

+//======================================================================================================================
+//
+//  This file is part of waLBerla. waLBerla is free software: you can
+//  redistribute it and/or modify it under the terms of the GNU General Public
+//  License as published by the Free Software Foundation, either version 3 of
+//  the License, or (at your option) any later version.
+//
+//  waLBerla is distributed in the hope that it will be useful, but WITHOUT
+//  ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or
+//  FITNESS FOR A PARTICULAR PURPOSE.  See the GNU General Public License
+//  for more details.
+//
+//  You should have received a copy of the GNU General Public License along
+//  with waLBerla (see COPYING.txt). If not, see <http://www.gnu.org/licenses/>.
+//
+//! \file StreamInCellIntervalCodegenTest.cpp
+//! \ingroup lbm/codegen
+//! \author Christoph Schwarzmeier <christoph.schwarzmeier@fau.de>
+//! \brief Test generated streamInCellInterval() by comparing with regular generated stream().
+//
+//! Initialize a periodic 2x1x1 PDF field with arbitrary values and stream the PDFs with both functions.
+//======================================================================================================================
+
+#include "blockforest/Initialization.h"
+#include "blockforest/communication/UniformBufferedScheme.h"
+
+#include "core/Environment.h"
+#include "core/debug/TestSubsystem.h"
+
+#include "lbm/communication/PdfFieldPackInfo.h"
+#include "lbm/field/Adaptors.h"
+#include "lbm/field/AddToStorage.h"
+#include "lbm/sweeps/CellwiseSweep.h"
+
+#include "timeloop/SweepTimeloop.h"
+
+#include "GeneratedLatticeModel.h"
+
+namespace walberla
+{
+namespace StreamInCellIntervalCodegenTest
+{
+using LatticeModel_T         = lbm::GeneratedLatticeModel;
+using Stencil_T              = LatticeModel_T::Stencil;
+using PdfField_T             = lbm::PdfField< LatticeModel_T >;
+using CommunicationStencil_T = LatticeModel_T::CommunicationStencil;
+
+int main(int argc, char** argv)
+{
+   debug::enterTestMode();
+   Environment env(argc, argv);
+
+   // define the domain size (2x1x1)
+   const Vector3< uint_t > numBlocks(uint_t(1));
+   const Vector3< uint_t > domainSize(uint_t(2), uint_t(1), uint_t(1));
+   const Vector3< uint_t > cellsPerBlock(domainSize[0] / numBlocks[0], domainSize[1] / numBlocks[1],
+                                         domainSize[2] / numBlocks[2]);
+
+   // ease access to both cells
+   const Cell cell0 = Cell(cell_idx_t(0), cell_idx_t(0), cell_idx_t(0));
+   const Cell cell1 = Cell(cell_idx_t(1), cell_idx_t(0), cell_idx_t(0));
+
+   // create block grid
+   const std::shared_ptr< StructuredBlockForest > blockForest =
+      blockforest::createUniformBlockGrid(numBlocks[0], numBlocks[1], numBlocks[2],             // blocks
+                                          cellsPerBlock[0], cellsPerBlock[1], cellsPerBlock[2], // cells
+                                          real_t(1.0),                                          // dx
+                                          true,                                                 // one block per process
+                                          true, true, true);                                    // periodicity
+
+   // relaxation rate
+   real_t omega = real_t(1.8);
+
+   // create lattice model
+   LatticeModel_T latticeModel = LatticeModel_T(omega);
+
+   // add PDF field (source for both functions, destination for regular stream)
+   const BlockDataID pdfFieldID =
+      lbm::addPdfFieldToStorage(blockForest, "PDF field", latticeModel, uint_c(1), field::fzyx);
+
+   // add PDF field as destination for streamInInterval()
+   const BlockDataID pdfFieldStreamIntervalID = lbm::addPdfFieldToStorage(
+      blockForest, "PDF destination field for streamInInterval", latticeModel, uint_c(1), field::fzyx);
+
+   // initialize PDF field with arbitrary values (should not be the same in both cells to improve validity of test)
+   for (auto blockIt = blockForest->begin(); blockIt != blockForest->end(); ++blockIt)
+   {
+      PdfField_T* pdfField = blockIt->getData< PdfField_T >(pdfFieldID);
+
+      pdfField->setDensityAndVelocity(cell0, Vector3< real_t >(real_t(0.1), real_t(0.1), real_t(0.1)), real_t(1.1));
+      pdfField->setDensityAndVelocity(cell1, Vector3< real_t >(real_t(-0.1), real_t(-0.1), real_t(-0.1)), real_t(0.9));
+   }
+
+   // create communication for PDF fields
+   blockforest::communication::UniformBufferedScheme< CommunicationStencil_T > communication(blockForest);
+   communication.addPackInfo(make_shared< lbm::PdfFieldPackInfo< LatticeModel_T > >(pdfFieldID));
+
+   // communicate PDF fields to fill ghost layers
+   communication();
+
+   for (auto blockIt = blockForest->begin(); blockIt != blockForest->end(); ++blockIt)
+   {
+      PdfField_T* pdfField               = blockIt->getData< PdfField_T >(pdfFieldID);
+      PdfField_T* pdfFieldStreamInterval = blockIt->getData< PdfField_T >(pdfFieldStreamIntervalID);
+
+      // use streamInCellInterval() (does not change content of pdfField)
+      auto lbmSweepGenerated = typename LatticeModel_T::Sweep(pdfFieldID);
+      const CellInterval ci  = CellInterval(cell0, cell1);
+      lbmSweepGenerated.streamInCellInterval(pdfField, pdfFieldStreamInterval, ci);
+
+      // use regular stream(); overwrites content of pdfField and MUST therefore be called after streamInCellInterval()
+      lbmSweepGenerated.stream(blockIt.get());
+   }
+
+   // check equality of streamInCellInterval() and stream()
+   for (auto blockIt = blockForest->begin(); blockIt != blockForest->end(); ++blockIt)
+   {
+      PdfField_T* pdfField               = blockIt->getData< PdfField_T >(pdfFieldID);
+      PdfField_T* pdfFieldStreamInterval = blockIt->getData< PdfField_T >(pdfFieldStreamIntervalID);
+
+      WALBERLA_FOR_ALL_CELLS(pdfFieldIt, pdfField, pdfFieldStreamIntervalIt, pdfFieldStreamInterval, {
+         // check equality of each PDF
+         for (uint_t i = uint_t(0); i != pdfField->F_SIZE; ++i)
+         {
+            WALBERLA_CHECK_FLOAT_EQUAL(pdfFieldIt[i], pdfFieldStreamIntervalIt[i]);
+         }
+      }) // WALBERLA_FOR_ALL_CELLS
+   }
+
+   return EXIT_SUCCESS;
+}
+} // namespace StreamInCellIntervalCodegenTest
+} // namespace walberla
+
+int main(int argc, char** argv) { return walberla::StreamInCellIntervalCodegenTest::main(argc, argv); }