Complete and partially tested code for boundary indexing

lbmpy/advanced_streaming/__init__.py

-from .indexing import AdvancedStreamingIndexing
+from .boundaryindexing import AdvancedStreamingBoundaryIndexing
 
-__all__ = ['AdvancedStreamingIndexing']
+__all__ = ['AdvancedStreamingBoundaryIndexing']

lbmpy/advanced_streaming/indexing.py → lbmpy/advanced_streaming/boundaryindexing.py

@@ -3,6 +3,7 @@ import sympy as sp
 import pystencils as ps
 
 from pystencils.data_types import TypedSymbol, create_type
+from pystencils.backends.cbackend import CustomCodeNode
 
 from lbmpy.fieldaccess import StreamPullTwoFieldsAccessor, \
     StreamPushTwoFieldsAccessor, \
@@ -11,8 +12,10 @@ from lbmpy.fieldaccess import StreamPullTwoFieldsAccessor, \
     EsoTwistEvenTimeStepAccessor, \
     EsoTwistOddTimeStepAccessor
 
+from itertools import product
 
-class AdvancedStreamingIndexing:
+
+class AdvancedStreamingBoundaryIndexing:
 
     #   =======================================
     #       Symbols for usage in boundaries
@@ -21,7 +24,7 @@ class AdvancedStreamingIndexing:
     def proxy_fields(self):
         q = len(self.stencil)
         d = len(self.stencil[0])
-        return ps.fields(f"f_out({q}), f_in({q}): [{d}]")
+        return ps.fields(f"f_out({q}), f_in({q}): [{d}D]")
 
     def dir_symbol(self):
         return TypedSymbol('dir', create_type(np.int64))
@@ -74,84 +77,136 @@ class AdvancedStreamingIndexing:
         inward_accesses = (
             even_accessor if odd_to_even else odd_accessor).read(pdf_field, stencil)
 
-        self.outward_accesses = outward_accesses
-        self.inward_accesses = inward_accesses
+        self.accesses = {'out': outward_accesses, 'in': inward_accesses}
 
         self.pdf_field = pdf_field
         self.stencil = stencil
-        self.directions = ['x', 'y', 'z'][:len(stencil[0])]
+        self.dim = len(stencil[0])
+        self.q = len(stencil)
+        self.coordinate_names = ['x', 'y', 'z'][:self.dim]
 
-        #   Collection of translation arrays required in generated code
         self.required_arrays = set()
+        self.trivial_translations = self._collect_trivial_translations()
 
-    def _index_array_symbol(self, f_dir, inverse, dir_symbol):
+    def _index_array_name(self, f_dir, inverse):
         assert f_dir in ['in', 'out']
         inv = '_inv' if inverse else ''
         name = f"f_{f_dir}{inv}_dir_idx"
-        self.required_arrays.add(name)
-        return sp.IndexedBase(name)[dir_symbol]
+        return name
 
-    def _offset_array_symbols(self, f_dir, inverse, dir_symbol):
+    def _offset_array_names(self, f_dir, inverse):
         assert f_dir in ['in', 'out']
         inv = '_inv' if inverse else ''
         name_base = f"f_{f_dir}{inv}_offsets_"
-        names = [name_base + d for d in self.directions]
-        self.required_arrays |= set(names)
-        return [sp.IndexedBase(n)[dir_symbol] for n in names]
+        names = [name_base + d for d in self.coordinate_names]
+        return names
+
+    def _array_symbols(self, f_dir, inverse, index):
+        if (f_dir, inverse) in self.trivial_translations:
+            return {'index': index, 'offsets': (0, ) * self.dim}
+
+        index_array_name = self._index_array_name(f_dir, inverse)
+        offset_array_names = self._offset_array_names(f_dir, inverse)
+        self.required_arrays.add((f_dir, inverse))
+        return {
+            'index': sp.IndexedBase(index_array_name)[index],
+            'offsets': tuple(sp.IndexedBase(n)[index] for n in offset_array_names)
+        }
 
     #   =================================
     #       Proxy fields substitution
     #   =================================
 
     def substitute_proxies(self, assignments):
-        outward_accesses = self.outward_accesses
-        inward_accesses = self.inward_accesses
+        if isinstance(assignments, ps.Assignment):
+            assignments = [assignments]
+
+        if not isinstance(assignments, ps.AssignmentCollection):
+            assignments = ps.AssignmentCollection(assignments)
 
+        accesses = self.accesses
         f_out, f_in = self.proxy_fields()
-        dir_symbol = self.dir_symbol()
         inv_dir = self.inverse_dir_symbol()
 
-        # 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] = self.pdf_field[
-                            self._offset_array_symbols('out', False, dir_symbol)
-                        ](self._index_array_symbol('out', False, dir_symbol))
-                    else:
-                        # other cases like inverse direction, etc.
-                        pass
-                else:
-                    # non-trivial neighbour accesses -> add neighbour offset to streaming pattern offsets
-                    pass
-
+                f_dir = 'out'
             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] = self.pdf_field[
-                            self._offset_array_symbols('in', True, dir_symbol)
-                        ](self._index_array_symbol('in', True, dir_symbol))
-                    else:
-                        # other cases
-                        pass
-                else:
-                    # non-trivial neighbour accesses -> add neighbour offset to streaming pattern offsets
-                    pass
-
+                f_dir = 'in'
             else:
-                pass
+                continue
+
+            inv = False
+            idx = fa.index[0]
+            if isinstance(idx, sp.Indexed) and idx.base == inv_dir:
+                idx = idx.indices[0]
+                if isinstance(sp.sympify(idx), sp.Integer):
+                    idx = self._inverse_integer_dir_index(idx)
+                inv = True
+
+            if isinstance(sp.sympify(idx), sp.Integer):
+                accessor_subs[fa] = accesses[f_dir][idx].get_shifted(*(fa.offsets))
+            else:
+                arr = self._array_symbols(f_dir, inv, idx)
+                accessor_subs[fa] = self.pdf_field[arr['offsets']](arr['index']).get_shifted(*(fa.offsets))
 
         return assignments.new_with_substitutions(accessor_subs)
 
-#   end class AdvancedStreamingIndexing
+    #   =================
+    #       Internals
+    #   =================
+
+    def _inverse_integer_dir_index(self, idx):
+        return self.stencil.index(tuple(-d for d in self.stencil[idx]))
+
+    def _get_translated_indices_and_offsets(self, f_dir, inv):
+        accesses = self.accesses[f_dir]
+
+        if inv:
+            inverse_indices = [self._inverse_integer_dir_index(i)
+                               for i in range(len(self.stencil))]
+            accesses = [accesses[idx] for idx in inverse_indices]
+
+        indices = [a.index[0] for a in accesses]
+        offsets = []
+        for d in range(self.dim):
+            offsets.append([a.offsets[d] for a in accesses])
+        return indices, offsets
+
+    def _collect_trivial_translations(self):
+        trivials = set()
+        trivial_indices = list(range(self.q))
+        trivial_offsets = [[0] * self.q] * self.dim
+        for f_dir, inv in product(['in', 'out'], [False, True]):
+            indices, offsets = self._get_translated_indices_and_offsets(f_dir, inv)
+            if indices == trivial_indices and offsets == trivial_offsets:
+                trivials.add((f_dir, inv))
+        return trivials
+
+    def code_node(self):
+        return AdvancedStreamingBoundaryIndexing.TranslationArraysNode(self)
+
+    class TranslationArraysNode(CustomCodeNode):
+
+        def __init__(self, indexing):
+            code = '\n'
+
+            for f_dir, inv in indexing.required_arrays:
+                indices, offsets = indexing._get_translated_indices_and_offsets(f_dir, inv)
+
+                acc_indices = ", ".join([str(i) for i in indices])
+                code += self._array_pattern(indexing._index_array_name(f_dir, inv), acc_indices)
+
+                for d, arrname in enumerate(indexing._offset_array_names(f_dir, inv)):
+                    acc_offsets = ", ".join([str(o) for o in offsets[d]])
+                    code += self._array_pattern(arrname, acc_offsets)
+
+            super(AdvancedStreamingBoundaryIndexing.TranslationArraysNode, self).__init__(
+                code, symbols_read=set(), symbols_defined=indexing.required_arrays)
+
+        def _array_pattern(self, name, content):
+            return f"const int64_t {name} [] = {{ {content} }}; \n"
+
+#   end class AdvancedStreamingIndexing
\ No newline at end of file

lbmpy_tests/test_flexible_noslip.py

@@ -2,7 +2,7 @@ import numpy as np
 import sympy as sp
 
 from lbmpy.fieldaccess import *
-from lbmpy.advanced_streaming import FlexibleNoSlip, create_flexible_lbm_boundary_kernel
+from lbmpy.advanced_streaming.flexible_boundaries_prototype import FlexibleNoSlip, create_flexible_lbm_boundary_kernel
 from lbmpy.creationfunctions import create_lb_method
 from lbmpy.stencils import get_stencil