diff --git a/pystencils/astnodes.py b/pystencils/astnodes.py index 393396aa2d073c77e8df6ff627860164c9cfe526..b24138286c024d0e7b546488289e552620a6d125 100644 --- a/pystencils/astnodes.py +++ b/pystencils/astnodes.py @@ -509,6 +509,13 @@ class SympyAssignment(Node): self.lhs = fast_subs(self.lhs, subs_dict) self.rhs = fast_subs(self.rhs, subs_dict) + def optimize(self, optimizations): + try: + from sympy.codegen.rewriting import optimize + self.rhs = optimize(self.rhs, optimizations) + except Exception: + pass + @property def args(self): return [self._lhs_symbol, self.rhs] diff --git a/pystencils/math_optimizations.py b/pystencils/math_optimizations.py new file mode 100644 index 0000000000000000000000000000000000000000..ad0114782e15977fb329baacfb6ad4b6bac1e33b --- /dev/null +++ b/pystencils/math_optimizations.py @@ -0,0 +1,46 @@ +""" +Default Sympy optimizations applied in pystencils kernels using :func:`sympy.codegen.rewriting.optimize`. + +See :func:`sympy.codegen.rewriting.optimize`. +""" + + +import itertools + +from pystencils import Assignment +from pystencils.astnodes import SympyAssignment + +try: + from sympy.codegen.rewriting import optims_c99, optimize + from sympy.codegen.rewriting import ReplaceOptim + HAS_REWRITING = True + + # Evaluates all constant terms + evaluate_constant_terms = ReplaceOptim( + lambda e: hasattr(e, 'is_constant') and e.is_constant and not e.is_integer, + lambda p: p.evalf() + ) + + optims_pystencils_cpu = [evaluate_constant_terms] + list(optims_c99) + optims_pystencils_gpu = [evaluate_constant_terms] + list(optims_c99) +except ImportError: + from warnings import warn + warn("Could not import ReplaceOptim, optims_c99, optimize from sympy.codegen.rewriting." + "Please update your sympy installation!") + optims_c99 = [] + optims_pystencils_cpu = [] + optims_pystencils_gpu = [] + HAS_REWRITING = False + + +def optimize_assignments(assignments, optimizations): + + if HAS_REWRITING: + assignments = [Assignment(a.lhs, optimize(a.rhs, optimizations)) + if hasattr(a, 'lhs') + else a for a in assignments] + assignments_nodes = [a.atoms(SympyAssignment) for a in assignments] + for a in itertools.chain.from_iterable(assignments_nodes): + a.optimize(optimizations) + + return assignments diff --git a/pystencils_tests/test_sympy_optimizations.py b/pystencils_tests/test_sympy_optimizations.py new file mode 100644 index 0000000000000000000000000000000000000000..745b936ea7d0e3ce5da6c684dee42c56c09dd835 --- /dev/null +++ b/pystencils_tests/test_sympy_optimizations.py @@ -0,0 +1,60 @@ +import pytest +import sympy as sp + +import pystencils +from pystencils.math_optimizations import HAS_REWRITING, optimize_assignments, optims_pystencils_cpu + + +@pytest.mark.skipif(not HAS_REWRITING, reason="need sympy.codegen.rewriting") +def test_sympy_optimizations(): + for target in ('cpu', 'gpu'): + x, y, z = pystencils.fields('x, y, z: float32[2d]') + + # Triggers Sympy's expm1 optimization + assignments = pystencils.AssignmentCollection({ + x[0, 0]: sp.exp(y[0, 0]) - 1 + }) + + assignments = optimize_assignments(assignments, optims_pystencils_cpu) + + ast = pystencils.create_kernel(assignments, target=target) + code = str(pystencils.show_code(ast)) + assert 'expm1(' in code + + +@pytest.mark.skipif(not HAS_REWRITING, reason="need sympy.codegen.rewriting") +def test_evaluate_constant_terms(): + for target in ('cpu', 'gpu'): + x, y, z = pystencils.fields('x, y, z: float32[2d]') + + # Triggers Sympy's expm1 optimization + assignments = pystencils.AssignmentCollection({ + x[0, 0]: -sp.cos(1) + y[0, 0] + }) + + assignments = optimize_assignments(assignments, optims_pystencils_cpu) + + ast = pystencils.create_kernel(assignments, target=target) + code = str(pystencils.show_code(ast)) + assert 'cos(' not in code + print(code) + + +@pytest.mark.skipif(not HAS_REWRITING, reason="need sympy.codegen.rewriting") +def test_do_not_evaluate_constant_terms(): + optimizations = pystencils.math_optimizations.optims_pystencils_cpu + optimizations.remove(pystencils.math_optimizations.evaluate_constant_terms) + + for target in ('cpu', 'gpu'): + x, y, z = pystencils.fields('x, y, z: float32[2d]') + + assignments = pystencils.AssignmentCollection({ + x[0, 0]: -sp.cos(1) + y[0, 0] + }) + + optimize_assignments(assignments, optimizations) + + ast = pystencils.create_kernel(assignments, target=target) + code = str(pystencils.show_code(ast)) + assert 'cos(' in code + print(code)