Commit 21122ee0 authored by Frederik Hennig's avatar Frederik Hennig
Browse files

Complete and partially tested code for boundary indexing

parent 8244c40e
Pipeline #26874 failed with stage
in 26 minutes and 58 seconds
from .indexing import AdvancedStreamingIndexing
from .boundaryindexing import AdvancedStreamingBoundaryIndexing
__all__ = ['AdvancedStreamingIndexing']
__all__ = ['AdvancedStreamingBoundaryIndexing']
......@@ -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
......@@ -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
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment