Merge branch 'sum-prod' into 'master'

Implement sp.Sum, sp.Product

See merge request pycodegen/pystencils!38

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:

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
 

pystencils_tests/test_sum_prod.py

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