From 5944a75930ee1729deec786616cdcb716c487316 Mon Sep 17 00:00:00 2001
From: Stephan Seitz <stephan.seitz@fau.de>
Date: Mon, 19 Aug 2019 19:56:58 +0200
Subject: [PATCH] Implement sp.Sum, sp.Product
Sum and Product have a indexing variable which is a Atom but not a free
symbol
---
pystencils/astnodes.py | 2 +-
pystencils/backends/cbackend.py | 45 ++++++++++
pystencils_tests/test_sum_prod.py | 132 ++++++++++++++++++++++++++++++
3 files changed, 178 insertions(+), 1 deletion(-)
create mode 100644 pystencils_tests/test_sum_prod.py
diff --git a/pystencils/astnodes.py b/pystencils/astnodes.py
index 461ced21b..b2f5727c4 100644
--- a/pystencils/astnodes.py
+++ b/pystencils/astnodes.py
@@ -521,7 +521,7 @@ class SympyAssignment(Node):
@property
def undefined_symbols(self):
- result = self.rhs.atoms(sp.Symbol)
+ result = {s for s in self.rhs.free_symbols if not isinstance(s, sp.Indexed)}
# Add loop counters if there a field accesses
loop_counters = set()
for symbol in result:
diff --git a/pystencils/backends/cbackend.py b/pystencils/backends/cbackend.py
index 0ca5e6295..cac2b00a1 100644
--- a/pystencils/backends/cbackend.py
+++ b/pystencils/backends/cbackend.py
@@ -371,6 +371,51 @@ class CustomSympyPrinter(CCodePrinter):
else:
return res
+ def _print_Sum(self, expr):
+ template = """[&]() {{
+ {dtype} sum = ({dtype}) 0;
+ for ( {iterator_dtype} {var} = {start}; {condition}; {var} += {increment} ) {{
+ sum += {expr};
+ }}
+ return sum;
+}}()"""
+ var = expr.limits[0][0]
+ start = expr.limits[0][1]
+ end = expr.limits[0][2]
+ code = template.format(
+ dtype=get_type_of_expression(expr.args[0]),
+ iterator_dtype='int',
+ var=self._print(var),
+ start=self._print(start),
+ end=self._print(end),
+ expr=self._print(expr.function),
+ increment=str(1),
+ condition=self._print(var) + ' <= ' + self._print(end) # if start < end else '>='
+ )
+ return code
+
+ def _print_Product(self, expr):
+ template = """[&]() {{
+ {dtype} product = ({dtype}) 1;
+ for ( {iterator_dtype} {var} = {start}; {condition}; {var} += {increment} ) {{
+ product *= {expr};
+ }}
+ return product;
+}}()"""
+ var = expr.limits[0][0]
+ start = expr.limits[0][1]
+ end = expr.limits[0][2]
+ code = template.format(
+ dtype=get_type_of_expression(expr.args[0]),
+ iterator_dtype='int',
+ var=self._print(var),
+ start=self._print(start),
+ end=self._print(end),
+ expr=self._print(expr.function),
+ increment=str(1),
+ condition=self._print(var) + ' <= ' + self._print(end) # if start < end else '>='
+ )
+ return code
_print_Max = C89CodePrinter._print_Max
_print_Min = C89CodePrinter._print_Min
diff --git a/pystencils_tests/test_sum_prod.py b/pystencils_tests/test_sum_prod.py
new file mode 100644
index 000000000..4fa5c0618
--- /dev/null
+++ b/pystencils_tests/test_sum_prod.py
@@ -0,0 +1,132 @@
+# -*- coding: utf-8 -*-
+#
+# Copyright © 2019 Stephan Seitz <stephan.seitz@fau.de>
+#
+# Distributed under terms of the GPLv3 license.
+
+"""
+
+"""
+import numpy as np
+import sympy
+from sympy.abc import k
+
+import pystencils
+from pystencils.data_types import create_type
+
+
+def test_sum():
+
+ sum = sympy.Sum(k, (k, 1, 100))
+ expanded_sum = sum.doit()
+
+ print(sum)
+ print(expanded_sum)
+
+ x = pystencils.fields('x: float32[1d]')
+
+ assignments = pystencils.AssignmentCollection({
+ x.center(): sum
+ })
+
+ ast = pystencils.create_kernel(assignments)
+ code = str(pystencils.show_code(ast))
+ kernel = ast.compile()
+
+ print(code)
+ assert 'double sum' in code
+
+ array = np.zeros((10,), np.float32)
+
+ kernel(x=array)
+
+ assert np.allclose(array, int(expanded_sum) * np.ones_like(array))
+
+
+def test_sum_use_float():
+
+ sum = sympy.Sum(k, (k, 1, 100))
+ expanded_sum = sum.doit()
+
+ print(sum)
+ print(expanded_sum)
+
+ x = pystencils.fields('x: float32[1d]')
+
+ assignments = pystencils.AssignmentCollection({
+ x.center(): sum
+ })
+
+ ast = pystencils.create_kernel(assignments, data_type=create_type('float32'))
+ code = str(pystencils.show_code(ast))
+ kernel = ast.compile()
+
+ print(code)
+ print(pystencils.show_code(ast))
+ assert 'float sum' in code
+
+ array = np.zeros((10,), np.float32)
+
+ kernel(x=array)
+
+ assert np.allclose(array, int(expanded_sum) * np.ones_like(array))
+
+
+def test_product():
+
+ k = pystencils.TypedSymbol('k', create_type('int64'))
+
+ sum = sympy.Product(k, (k, 1, 10))
+ expanded_sum = sum.doit()
+
+ print(sum)
+ print(expanded_sum)
+
+ x = pystencils.fields('x: int64[1d]')
+
+ assignments = pystencils.AssignmentCollection({
+ x.center(): sum
+ })
+
+ ast = pystencils.create_kernel(assignments)
+ code = str(pystencils.show_code(ast))
+ kernel = ast.compile()
+
+ print(code)
+ assert 'int64_t product' in code
+
+ array = np.zeros((10,), np.int64)
+
+ kernel(x=array)
+
+ assert np.allclose(array, int(expanded_sum) * np.ones_like(array))
+
+
+def test_prod_var_limit():
+
+ k = pystencils.TypedSymbol('k', create_type('int64'))
+ limit = pystencils.TypedSymbol('limit', create_type('int64'))
+
+ sum = sympy.Sum(k, (k, 1, limit))
+ expanded_sum = sum.replace(limit, 100).doit()
+
+ print(sum)
+ print(expanded_sum)
+
+ x = pystencils.fields('x: int64[1d]')
+
+ assignments = pystencils.AssignmentCollection({
+ x.center(): sum
+ })
+
+ ast = pystencils.create_kernel(assignments)
+ code = str(pystencils.show_code(ast))
+ kernel = ast.compile()
+
+ print(code)
+
+ array = np.zeros((10,), np.int64)
+
+ kernel(x=array, limit=100)
+
+ assert np.allclose(array, int(expanded_sum) * np.ones_like(array))
--
GitLab