From 8479aa27486f408c7e0d9099d966d79b05ffd205 Mon Sep 17 00:00:00 2001
From: Stephan Seitz <stephan.seitz@fau.de>
Date: Sun, 18 Aug 2019 11:23:50 +0200
Subject: [PATCH 1/4] Let iterator of AssignmentCollection iterate over
all_assignments (not only main_assignments)
---
pystencils/simp/assignment_collection.py | 2 +-
1 file changed, 1 insertion(+), 1 deletion(-)
diff --git a/pystencils/simp/assignment_collection.py b/pystencils/simp/assignment_collection.py
index fa437a5d0..1e9a4c47e 100644
--- a/pystencils/simp/assignment_collection.py
+++ b/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):
--
GitLab
From da5cfa52ef579ffac74ba635bc04a34d8365d6e6 Mon Sep 17 00:00:00 2001
From: Stephan Seitz <stephan.seitz@fau.de>
Date: Mon, 23 Sep 2019 09:20:44 +0200
Subject: [PATCH 2/4] Add dtype assumptions to cast_func
---
pystencils/data_types.py | 31 +++++++++++++++++++++++++++++++
1 file changed, 31 insertions(+)
diff --git a/pystencils/data_types.py b/pystencils/data_types.py
index 86ce1747d..faa886718 100644
--- a/pystencils/data_types.py
+++ b/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):
--
GitLab
From c3a4fb73647d84b02a4d4942f2a5b883d574cf8e Mon Sep 17 00:00:00 2001
From: Stephan Seitz <stephan.seitz@fau.de>
Date: Mon, 23 Sep 2019 09:23:57 +0200
Subject: [PATCH 3/4] Add property is_integer -> True to
IntegerFunctionTwoArgsMixIn
---
pystencils/integer_functions.py | 4 ++++
1 file changed, 4 insertions(+)
diff --git a/pystencils/integer_functions.py b/pystencils/integer_functions.py
index b54bbaab2..fa5a4d739 100644
--- a/pystencils/integer_functions.py
+++ b/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):
--
GitLab
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 4/4] 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 393396aa2..b24138286 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 000000000..ad0114782
--- /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 000000000..745b936ea
--- /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