diff --git a/pystencils/astnodes.py b/pystencils/astnodes.py index 461ced21b013b442c55589ff668a3cbb65d8fee6..b2f5727c4068852fd6f4ba1b755a1ba5d1215b91 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 0ca5e6295cfd14177fba63830ed3dd1fe95e28bd..cac2b00a1b52b29a61ccaa0cd73e9e5d37a89e95 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 0000000000000000000000000000000000000000..4fa5c0618612b013edda4d164dd035dafdd2438a --- /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))