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

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

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

See merge request pycodegen/pystencils!34
parents 71b8767b 8d610054
......@@ -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]
......
......@@ -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):
......
......@@ -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):
......
"""
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
......@@ -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):
......
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)
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