Commit 5c876796 authored by Frederik Hennig's avatar Frederik Hennig
Browse files

Advanced Streaming Indexing WIP

parent 27e9d65b
Pipeline #26834 failed with stage
in 6 minutes and 17 seconds
from .flexible_boundaries import AdvancedStreamingBoundaryOffsetInfo, \
FlexibleNoSlip, create_flexible_lbm_boundary_kernel
from .indexing import AdvancedStreamingIndexing
__all__ = [ 'AdvancedStreamingBoundaryOffsetInfo',
'create_flexible_lbm_boundary_kernel' ]
\ No newline at end of file
__all__ = [ 'AdvancedStreamingIndexing' ]
\ No newline at end of file
import numpy as np
import sympy as sp
import pystencils as ps
from pystencils.data_types import TypedSymbol, create_type
from lbmpy.fieldaccess import StreamPullTwoFieldsAccessor, \
StreamPushTwoFieldsAccessor, \
AAEvenTimeStepAccessor, \
AAOddTimeStepAccessor, \
EsoTwistEvenTimeStepAccessor, \
class AdvancedStreamingIndexing:
# =======================================
# Symbols for usage in boundaries
# =======================================
def proxy_fields(self):
q = len(self.stencil)
d = len(self.stencil[0])
return ps.fields('f_out({q}), f_in({q}): [{d}}]'.format(q=q, d=d))
def dir_symbol(self):
return TypedSymbol('dir', create_type(np.int64))
def inverse_dir_symbol(self):
return sp.IndexedBase('invdir')
# =============================
# Constructor and State
# =============================
def __init__(self, pdf_field, stencil, between_timesteps = 'both', kernel_type='pull'):
if between_timesteps not in ['both', 'odd_to_even', 'even_to_odd']:
raise ValueError(
"Invalid value of parameter 'between_timesteps'. Must be one of 'both', 'odd_to_even' or 'even_to_odd'.",
if between_timesteps == 'both' and kernel_type in ['aa', 'esotwist']:
raise ValueError('Cannot create an offset info for both kinds of timesteps for kernel type ' + kernel_type)
odd_to_even = (between_timesteps == 'odd_to_even')
even_accessors = {
'pull': StreamPullTwoFieldsAccessor,
'push': StreamPushTwoFieldsAccessor,
'aa': AAEvenTimeStepAccessor,
'esotwist': EsoTwistEvenTimeStepAccessor
odd_accessors = {
'pull': StreamPullTwoFieldsAccessor,
'push': StreamPushTwoFieldsAccessor,
'aa': AAOddTimeStepAccessor,
'esotwist': EsoTwistOddTimeStepAccessor
even_accessor = even_accessors[kernel_type]
odd_accessor = odd_accessors[kernel_type]
except KeyError:
raise ValueError(
"Invalid value of parameter 'kernel_type'.", kernel_type)
if between_timesteps == 'both':
assert even_accessor == odd_accessor
outward_accesses = (
odd_accessor if odd_to_even else even_accessor).write(pdf_field, stencil)
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.pdf_field = pdf_field
self.stencil = stencil
self.directions = ['x', 'y', 'z'][:len(stencil[0])]
# Collection of translation arrays required in generated code
self.required_arrays = set()
def _index_array_symbol(self, f_dir, inverse, dir_symbol):
assert f_dir in ['in', 'out']
name = "f_{d}{inv}_dir_idx".format(d=f_dir, inv='_inv' if inverse else '')
return sp.IndexedBase(name)[dir_symbol]
def _offset_array_symbols(self, f_dir, inverse, dir_symbol):
assert f_dir in ['in', 'out']
name_base = "f_{d}{inv}_offsets_".format(d=f_dir, inv='_inv' if inverse else '')
names = [ name_base + d for d in self.directions ]
self.required_arrays |= set(names)
return [ sp.IndexedBase(n)[dir_symbol] for n in names ]
# =================================
# Proxy fields substitution
# =================================
def substitute_proxies(self, assignments):
# Get the accessor lists for the streaming pattern (here AA)
outward_accesses = self.outward_accesses
inward_accesses = self.inward_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))
# other cases like inverse direction, etc.
# non-trivial neighbour accesses -> add neighbour offset to streaming pattern offsets
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))
# other cases
# non-trivial neighbour accesses -> add neighbour offset to streaming pattern offsets
return assignments.new_with_substitutions(accessor_subs)
# end class AdvancedStreamingIndexing
\ No newline at end of file
Supports Markdown
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