Commit d5373ec8 authored by Martin Bauer's avatar Martin Bauer
Browse files

Fix for LB collide only step

- sweep updates pdf field in place, but does not load all the values
  before overwriting them
-> added a new transformation on AssignmentCollection that loads all
   read values first
parent 1b1ebcd7
...@@ -2,9 +2,10 @@ from .assignment_collection import AssignmentCollection ...@@ -2,9 +2,10 @@ from .assignment_collection import AssignmentCollection
from .simplificationstrategy import SimplificationStrategy from .simplificationstrategy import SimplificationStrategy
from .simplifications import sympy_cse, sympy_cse_on_assignment_list, \ from .simplifications import sympy_cse, sympy_cse_on_assignment_list, \
apply_to_all_assignments, apply_on_all_subexpressions, subexpression_substitution_in_existing_subexpressions, \ apply_to_all_assignments, apply_on_all_subexpressions, subexpression_substitution_in_existing_subexpressions, \
subexpression_substitution_in_main_assignments, add_subexpressions_for_divisions subexpression_substitution_in_main_assignments, add_subexpressions_for_divisions, add_subexpressions_for_field_reads
__all__ = ['AssignmentCollection', 'SimplificationStrategy', __all__ = ['AssignmentCollection', 'SimplificationStrategy',
'sympy_cse', 'sympy_cse_on_assignment_list', 'apply_to_all_assignments', 'sympy_cse', 'sympy_cse_on_assignment_list', 'apply_to_all_assignments',
'apply_on_all_subexpressions', 'subexpression_substitution_in_existing_subexpressions', 'apply_on_all_subexpressions', 'subexpression_substitution_in_existing_subexpressions',
'subexpression_substitution_in_main_assignments', 'add_subexpressions_for_divisions'] 'subexpression_substitution_in_main_assignments', 'add_subexpressions_for_divisions',
'add_subexpressions_for_field_reads']
import sympy as sp import sympy as sp
from typing import Callable, List from typing import Callable, List
from pystencils import Field
from pystencils.assignment import Assignment from pystencils.assignment import Assignment
from pystencils.simp.assignment_collection import AssignmentCollection from pystencils.simp.assignment_collection import AssignmentCollection
from pystencils.sympyextensions import subs_additive from pystencils.sympyextensions import subs_additive
...@@ -83,6 +85,24 @@ def add_subexpressions_for_divisions(ac: AC) -> AC: ...@@ -83,6 +85,24 @@ def add_subexpressions_for_divisions(ac: AC) -> AC:
return ac.new_with_substitutions(substitutions, True) return ac.new_with_substitutions(substitutions, True)
def add_subexpressions_for_field_reads(ac: AC, subexpressions=True, main_assignments=True) -> AC:
r"""Substitutes field accesses on rhs of assignments with subexpressions
Can change semantics of the update rule (which is the goal of this transformation)
This is useful if a field should be update in place - all values are loaded before into subexpression variables,
then the new values are computed and written to the same field in-place.
"""
field_reads = set()
if subexpressions:
for assignment in ac.subexpressions:
field_reads.update(assignment.rhs.atoms(Field.Access))
if main_assignments:
for assignment in ac.main_assignments:
field_reads.update(assignment.rhs.atoms(Field.Access))
substitutions = {fa: sp.Dummy() for fa in field_reads}
return ac.new_with_substitutions(substitutions, add_substitutions_as_subexpressions=True, substitute_on_lhs=False)
def apply_to_all_assignments(operation: Callable[[sp.Expr], sp.Expr]) -> Callable[[AC], AC]: def apply_to_all_assignments(operation: Callable[[sp.Expr], sp.Expr]) -> Callable[[AC], AC]:
"""Applies sympy expand operation to all equations in collection.""" """Applies sympy expand operation to all equations in collection."""
def f(assignment_collection: AC) -> AC: def f(assignment_collection: AC) -> AC:
......
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