From 8d6100541dddc5b0bbe5fc40361fd99566d0c13f Mon Sep 17 00:00:00 2001 From: Stephan Seitz <stephan.seitz@fau.de> Date: Mon, 23 Sep 2019 09:28:04 +0200 Subject: [PATCH] Add pystencils.math_optimizations --- pystencils/astnodes.py | 7 +++ pystencils/math_optimizations.py | 46 +++++++++++++++ pystencils_tests/test_sympy_optimizations.py | 60 ++++++++++++++++++++ 3 files changed, 113 insertions(+) create mode 100644 pystencils/math_optimizations.py create mode 100644 pystencils_tests/test_sympy_optimizations.py diff --git a/pystencils/astnodes.py b/pystencils/astnodes.py index 393396aa..b2413828 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 00000000..ad011478 --- /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 00000000..745b936e --- /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) -- GitLab