Merge branch 'use_codegen.rewriting.optimize' into 'master'

Address #13: Use sympy.codegen.rewriting.optimize

See merge request pycodegen/pystencils!34

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]

pystencils/data_types.py

@@ -82,6 +82,37 @@ class cast_func(sp.Function):
     def dtype(self):
         return self.args[1]
 
+    @property
+    def is_integer(self):
+        if hasattr(self.dtype, 'numpy_dtype'):
+            return np.issubdtype(self.dtype.numpy_dtype, np.integer) or super().is_integer
+        else:
+            return super().is_integer
+
+    @property
+    def is_negative(self):
+        if hasattr(self.dtype, 'numpy_dtype'):
+            if np.issubdtype(self.dtype.numpy_dtype, np.unsignedinteger):
+                return False
+
+        return super().is_negative
+
+    @property
+    def is_nonnegative(self):
+        if self.is_negative is False:
+            return True
+        else:
+            return super().is_nonnegative
+
+    @property
+    def is_real(self):
+        if hasattr(self.dtype, 'numpy_dtype'):
+            return np.issubdtype(self.dtype.numpy_dtype, np.integer) or \
+                np.issubdtype(self.dtype.numpy_dtype, np.floating) or \
+                super().is_real
+        else:
+            return super().is_real
+
 
 # noinspection PyPep8Naming
 class boolean_cast_func(cast_func, Boolean):

pystencils/integer_functions.py

@@ -25,6 +25,10 @@ class IntegerFunctionTwoArgsMixIn(sp.Function):
                 raise ValueError("Integer functions can only be constructed with typed expressions")
         return super().__new__(cls, *args)
 
+    @property
+    def is_integer(self):
+        return True
+
 
 # noinspection PyPep8Naming
 class bitwise_xor(IntegerFunctionTwoArgsMixIn):

pystencils/math_optimizations.py

+"""
+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

pystencils/simp/assignment_collection.py

@@ -347,7 +347,7 @@ class AssignmentCollection:
         return result
 
     def __iter__(self):
-        return self.main_assignments.__iter__()
+        return self.all_assignments.__iter__()
 
     @property
     def main_assignments_dict(self):

pystencils_tests/test_sympy_optimizations.py

+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)