Advanced Streaming Extensions

This patch extends lbmpy's support for advanced streaming patterns. In particular, the Push-Collide-Stream, AA-Pattern and Esoteric Twist are considered.

• Boundary Handling: The LBM boundary implementations are overhauled completely. A new framework for defining and generating boundary kernels is introduced, which makes the symbolic definition independent of the streaming pattern. Boundaries can now be defined symbolically using the proxy fields f_out for accessing populations streaming out of a cell, and f_in for populations streaming in to a cell. Those are replaced during the code generation process according to the streaming pattern. The class BetweenTimestepsIndexing used herein can also be used to generate other kernels that run on the PDF-field in-between collision steps.

• Periodicity Handling: The build-in periodicity synchronization of pystencil's DataHandling does not work for advanced LBM streaming. The new class LBMPeriodicityHandling extends the SerialDataHandling to handle periodicity correctly.

• Utility: The patch introduces a few utility functions for the work with advanced streaming patterns.

lbmpy/advanced_streaming/Advanced Boundaries.ipynb

+%% Cell type:code id: tags:
+``` 
+%load_ext autoreload
+%autoreload 2
+%% Cell type:code id: tags:
+``` 
+import sympy as sp
+import numpy as np
+import pystencils as ps
+
+from lbmpy.stencils import get_stencil
+from lbmpy.boundaries import NoSlip
+from lbmpy.boundaries.boundaryhandling import create_lattice_boltzmann_boundary_kernel
+from lbmpy.creationfunctions import create_lb_method
+
+stencil = get_stencil('D2Q9')
+pdf_field = ps.fields('pdf_field(9): [2D]')
+
+from pystencils.boundaries.createindexlist import numpy_data_type_for_boundary_object
+from pystencils.data_types import TypedSymbol, create_type
+from pystencils.field import FieldType
+
+lb_method = create_lb_method(stencil=stencil)
+noslip = NoSlip()
+index_struct_dtype = numpy_data_type_for_boundary_object(noslip, lb_method.dim)
+
+index_field = ps.Field('indexVector', FieldType.INDEXED, index_struct_dtype, layout=[0],
+                    shape=(TypedSymbol("indexVectorSize", create_type(np.int64)), 1), strides=(1, 1))
+%% Cell type:markdown id: tags:
+## Defining Boundaries: Current Implementaiton
+
+The boundary cell is accessed directly from the fluid cell through the `neighbour` offset with the inverse direction `inverse_dir`. Both of those are symbolic and are replaced by array accesses during code generation.
+```
+class NoSlip:
+    def __call__(self, pdf_field, direction_symbol, lb_method, **kwargs):
+        neighbor = BoundaryOffsetInfo.offset_from_dir(direction_symbol, lb_method.dim)
+        inverse_dir = BoundaryOffsetInfo.inv_dir(direction_symbol)
+        return [Assignment(pdf_field[neighbor](inverse_dir), pdf_field(direction_symbol))]
+```
+%% Cell type:code id: tags:
+``` 
+dir_symbol = TypedSymbol('dir', create_type(np.int64))
+noslip(pdf_field, dir_symbol, lb_method)[0]
+%% Output
+$\displaystyle {{pdf_field}_{({cx}_{dir},{cy}_{dir})}^{invdir[dir]}} \leftarrow {{pdf_field}_{(0,0)}^{dir}}$
+Assignment(pdf_field_dd4358bcc3a3, pdf_field_8ffb2303630e)
+%% Cell type:code id: tags:
+``` 
+noslip_ast = create_lattice_boltzmann_boundary_kernel(pdf_field, index_field, lb_method, noslip)
+ps.show_code(noslip_ast)
+%% Output
+%% Cell type:markdown id: tags:
+## Defining Boundaries: Flexible Approach
+
+The arrangment of the populations in memory is abstracted away through two proxy fields
+$$
+    f_{out}, \quad f_{in}
+$$
+for the populations streaming *out* of the current cell and *into* the current cell. With those, boundaries can be defined in general notation, without knowledge about the streaming pattern. For example, the NoSlip boundary is defined as:
+%% Cell type:code id: tags:
+``` 
+class FlexibleNoSlip:
+    def __call__(self, f_out, f_in, dir_symbol, inv_dir):
+        return ps.Assignment(f_in(inv_dir[dir_symbol]), f_out(dir_symbol))
+%% Cell type:code id: tags:
+``` 
+f_out, f_in = ps.fields('f_out(9), f_in(9): [2D]')
+dir_symbol = TypedSymbol('dir', create_type(np.int64))
+inv_dir = sp.IndexedBase('invdir')
+
+flex_noslip = FlexibleNoSlip()
+noslip_assignments =  flex_noslip(f_out, f_in, dir_symbol, inv_dir)
+noslip_assignments
+%% Output
+$\displaystyle {{f_in}_{(0,0)}^{invdir[dir]}} \leftarrow {{f_out}_{(0,0)}^{dir}}$
+Assignment(f_in_1ba4ca851771, f_out_8ffb2303630e)
+%% Cell type:markdown id: tags:
+During Code Generation, all accesses to $f_{out}$ and $f_{in}$ are replaced by accesses to `pdf_field` according to the streaming pattern.
+
+### Example: Boundary handling after an even time step with the AA-Pattern
+
+For the current cell, the populations flowing out are the ones written in the previous step, and the populations flowing in are those begin read in the next step. The previous step was controlled by the `AAEvenTimeStepAccessor`, and the next step is controlled by `AAOddTimeStepAccessor`. The required memory locations are thus given by:
+
+- Outbound populations: `AAEvenTimeStepAccessor.write`
+- Inbound populations: `AAOddTimeStepAccessor.read`
+%% Cell type:code id: tags:
+``` 
+from lbmpy.fieldaccess import AAEvenTimeStepAccessor, AAOddTimeStepAccessor
+
+def substitute_proxies(assignments):
+    # Get the accessor lists for the streaming pattern (here AA)
+    outward_accesses = AAEvenTimeStepAccessor.write(pdf_field, stencil)
+    inward_accesses = AAOddTimeStepAccessor.read(pdf_field, stencil)
+
+    # Define the symbols which will correspond to the index translation arrays
+    # in the generated code
+    f_out_dir_idx = sp.IndexedBase('f_out_dir_idx')
+    f_out_offsets_x = sp.IndexedBase('f_out_offset_x')
+    f_out_offsets_y = sp.IndexedBase('f_out_offset_y')
+
+    f_in_inv_dir_idx = sp.IndexedBase('f_in_inv_dir_idx')
+    f_in_inv_offsets_x = sp.IndexedBase('f_in_inv_offsets_x')
+    f_in_inv_offsets_y = sp.IndexedBase('f_in_inv_offsets_y')
+
+    # Substitute all proxy field accesses
+
+    if not isinstance(assignments, ps.AssignmentCollection):
+        assignments = ps.AssignmentCollection([assignments])
+
+    accessor_subs = dict()
+
+    for fa in assignments.atoms(ps.Field.Access):
+        if fa.field == f_out:
+            if fa.offsets == (0,0):
+                if isinstance(fa.index[0], int):
+                    accessor_subs[fa] = outward_accesses[fa.index[0]]
+                elif fa.index[0] == dir_symbol:
+                    accessor_subs[fa] = pdf_field[
+                            f_out_offsets_x[dir_symbol], f_out_offsets_y[dir_symbol]
+                        ]( f_out_dir_idx[dir_symbol] )
+                else:
+                    # other cases like inverse direction, etc.
+                    pass
+            else:
+                # non-trivial neighbour accesses -> add neighbour offset to streaming pattern offsets
+                pass
+
+        elif fa.field == f_in:
+            if fa.offsets == (0,0):
+                if isinstance(fa.index[0], int):
+                    accessor_subs[fa] = inward_accesses[fa.index[0]]
+                elif fa.index[0] == inv_dir[dir_symbol]:
+                    accessor_subs[fa] = pdf_field[
+                            f_in_inv_offsets_x[dir_symbol], f_in_inv_offsets_y[dir_symbol]
+                        ]( f_in_inv_dir_idx[dir_symbol] )
+                else:
+                    # other cases
+                    pass
+            else:
+                # non-trivial neighbour accesses -> add neighbour offset to streaming pattern offsets
+                pass
+
+        else:
+            pass
+
+    return assignments.new_with_substitutions(accessor_subs)
+
+
+substitute_proxies(noslip_assignments).main_assignments[0]
+%% Output
+$\displaystyle {{pdf_field}_{({f_{in inv offsets x}}_{dir},{f_{in inv offsets y}}_{dir})}^{f_in_inv_dir_idx[dir]}} \leftarrow {{pdf_field}_{({f_{out offset x}}_{dir},{f_{out offset y}}_{dir})}^{f_out_dir_idx[dir]}}$
+Assignment(pdf_field_5ed29f4dfe26, pdf_field_9520cd7bd10d)
+%% Cell type:markdown id: tags:
+Accesses with fixed indices are replaced directly:
+%% Cell type:code id: tags:
+``` 
+as2 = ps.Assignment(f_in(3), f_out(4) + f_out(5) + f_out(dir_symbol))
+as2
+%% Output
+$\displaystyle {{f_in}_{(0,0)}^{3}} \leftarrow {{f_out}_{(0,0)}^{dir}} + {{f_out}_{(0,0)}^{4}} + {{f_out}_{(0,0)}^{5}}$
+Assignment(f_in_C^3, f_out_8ffb2303630e + f_out_C^4 + f_out_C^5)
+%% Cell type:code id: tags:
+``` 
+substitute_proxies(as2).main_assignments[0]
+%% Output
+$\displaystyle {{pdf_field}_{(1,0)}^{4}} \leftarrow {{pdf_field}_{({f_{out offset x}}_{dir},{f_{out offset y}}_{dir})}^{f_out_dir_idx[dir]}} + {{pdf_field}_{(0,0)}^{3}} + {{pdf_field}_{(0,0)}^{8}}$
+Assignment(pdf_field_E^4, pdf_field_9520cd7bd10d + pdf_field_C^3 + pdf_field_C^8)
+%% Cell type:markdown id: tags:
+The index translation arrays are prepended in the generated code:
+```
+    const int64_t f_out_dir_idx [] = { 0, 2, 1, 4, 3, 8, 7, 6, 5 };
+    const int64_t f_out_offsets_x [] = { 0, 0, 0, 0, 0, 0, 0, 0, 0 };
+    const int64_t f_out_offsets_y [] = { 0, 0, 0, 0, 0, 0, 0, 0, 0 };
+    const int64_t f_in_inv_dir_idx [] = { 0, 1, 2, 3, 4, 5, 6, 7, 8 };
+    const int64_t f_in_inv_offsets_x [] = { 0, 0, 0, -1, 1, -1, 1, -1, 1 };
+    const int64_t f_in_inv_offsets_y [] = { 0, 1, -1, 0, 0, 1, 1, -1, -1 };
+```
+
+The generated boundary kernel is of course only fit for application after an even time step. For the AA-pattern's odd time steps, a second kernel needs to be generated. The definition of the NoSlip boundary itself is the same; only the index translation arrays are different.
+
+## Integration with the pystencils Boundary Handling Infrastructure
+
+The existing `pystencils.boundaries.BoundaryHandling` class is extended by a new subclass `FlexibleLBMBoundaryHandling` analogously to the implementation of the existing `LatticeBoltzmannBoundaryHandling`. Its behaviour depends on the streaming pattern which is passed to its constructor. It overrides `__call__` and `_create_boundary_kernel`:
+
+```
+class FlexibleLBMBoundaryHandling(BoundaryHandling):
+
+    def __init__(self, kernel_type): # kernel_type in ['pull', 'push', 'aa', 'esotwist']
+        (...)
+
+    def __call__(self, **kwargs):
+        # Handle boundaries like in base class, but call different kernels
+        # depending on the time step modulus
+
+    def _create_boundary_kernel(self, **kwargs):
+        if self.kernel_type in ['aa', 'esotwist']:
+            # in-place methods: create two boundary kernels, for even and odd timesteps
+        else:
+            # two-fields methods: create only one boundary kernel
+```
+
+## Integration with waLBerla
+
+The waLBerla boundary code generation can also be extended to generate implementations which track the time step internally and selectively call one of two implementations:
+
+```
+    class BoundarySweep{
+        public:
+            void operator() (IBlock * block){
+                if(even_timestep) evenSweep(block);
+                else oddSweep(block);
+            }
+    }
+```