diff --git a/.gitlab-ci.yml b/.gitlab-ci.yml index d7e410860a81170026ddde59c01bca42f2793bac..c4e477014fba2f550356d91c711b0fcedaa78db9 100644 --- a/.gitlab-ci.yml +++ b/.gitlab-ci.yml @@ -142,17 +142,13 @@ ubuntu: - sed -i 's/--doctest-modules //g' pytest.ini - env - pip3 list - - pytest-3 -v -n $NUM_CORES --cov-report html --cov-report term --cov=. --junitxml=report.xml pystencils_tests/test_*vec*.py pystencils_tests/test_random.py - - python3 -m coverage xml + - pytest-3 -v -n $NUM_CORES --junitxml=report.xml pystencils_tests/test_*vec*.py pystencils_tests/test_random.py tags: - docker - AVX artifacts: when: always - paths: - - coverage_report reports: - cobertura: coverage.xml junit: report.xml arm64v8: @@ -252,6 +248,7 @@ pycodegen-integration: - pip install -e pystencils/ - pip install -e lbmpy/ - pip install -e pygrandchem/ + - cmake --version - ./install_walberla.sh - export NUM_CORES=$(nproc --all) - mkdir -p ~/.config/matplotlib diff --git a/doc/index.rst b/doc/index.rst index a52161449e714d431cc258a06503398a986c8bf1..3030d33618303ecaf3c3d3f97e2fa12b73177965 100644 --- a/doc/index.rst +++ b/doc/index.rst @@ -14,5 +14,6 @@ pystencils can help you to generate blazingly fast code for image processing, nu .. image:: /img/pystencils_arch_block_diagram.svg :height: 450px + :align: center diff --git a/doc/sphinx/simplifications.rst b/doc/sphinx/simplifications.rst index ca1508c5b3a2b458d039977c735cf17003b692b1..63f394ca5d128ca6b70a9ebac4836cb7213aa21a 100644 --- a/doc/sphinx/simplifications.rst +++ b/doc/sphinx/simplifications.rst @@ -10,13 +10,27 @@ AssignmentCollection :members: +SimplificationStrategy +====================== + +.. autoclass:: pystencils.simp.SimplificationStrategy + :members: + Simplifications =============== -.. automodule:: pystencils.simp - :members: +.. automodule:: pystencils.simp.simplifications + :members: + +Subexpression insertion +======================= +The subexpression insertions have the goal to insert subexpressions which will not reduce the number of FLOPs. +For example a constant value kept as subexpression will lead to a new variable in the code which will occupy +a register slot. On the other side a single variable could just be inserted in all assignments. +.. automodule:: pystencils.simp.subexpression_insertion + :members: diff --git a/pystencils/cpu/cpujit.py b/pystencils/cpu/cpujit.py index ac1e8ac4f018b18ee331576ce87490b48d0ad6d1..ff577c442bd4828e2f29cf69d22a4466a5039f1a 100644 --- a/pystencils/cpu/cpujit.py +++ b/pystencils/cpu/cpujit.py @@ -554,7 +554,7 @@ class ExtensionModuleCode: if os.path.exists(os.path.join(os.path.dirname(__file__), '..', 'include', h[1:-1]))] header_hash = b''.join([hashlib.sha256(open(h, 'rb').read()).digest() for h in ps_headers]) - includes = "\n".join(["#include %s" % (include_file,) for include_file in header_list]) + includes = "\n".join([f"#include {include_file}" for include_file in header_list]) self._code_string += includes self._code_string += "\n" self._code_string += f"#define RESTRICT {restrict_qualifier} \n" @@ -563,7 +563,7 @@ class ExtensionModuleCode: for ast, name in zip(self._ast_nodes, self._function_names): old_name = ast.function_name - ast.function_name = "kernel_" + name + ast.function_name = f"kernel_{name}" self._code_string += generate_c(ast, custom_backend=self._custom_backend) self._code_string += create_function_boilerplate_code(ast.get_parameters(), name, ast) ast.function_name = old_name diff --git a/pystencils/cpu/kernelcreation.py b/pystencils/cpu/kernelcreation.py index 4237dcd34316ba75d8b83838c4e1df804e5b2efb..865beefa9b793233ea74c5b02b315371dcd6ce8e 100644 --- a/pystencils/cpu/kernelcreation.py +++ b/pystencils/cpu/kernelcreation.py @@ -205,5 +205,5 @@ def add_openmp(ast_node, schedule="static", num_threads=True, collapse=None, ass prefix = f"#pragma omp for schedule({schedule})" if collapse: - prefix += " collapse(%d)" % (collapse, ) + prefix += f" collapse({collapse})" loop_to_parallelize.prefix_lines.append(prefix) diff --git a/pystencils/kernelcreation.py b/pystencils/kernelcreation.py index e50851ae5a74676cfb71018cdd1a64da38b73645..ac4412256fa3735a003da50eaba114bdbda822a4 100644 --- a/pystencils/kernelcreation.py +++ b/pystencils/kernelcreation.py @@ -14,6 +14,8 @@ from pystencils.enums import Target, Backend from pystencils.field import Field, FieldType from pystencils.gpucuda.indexing import indexing_creator_from_params from pystencils.simp.assignment_collection import AssignmentCollection +from pystencils.simp.simplifications import apply_sympy_optimisations +from pystencils.simplificationfactory import create_simplification_strategy from pystencils.stencil import direction_string_to_offset, inverse_direction_string from pystencils.transformations import ( loop_blocking, move_constants_before_loop, remove_conditionals_in_staggered_kernel) @@ -83,6 +85,18 @@ class CreateKernelConfig: Dict with indexing parameters (constructor parameters of indexing class) e.g. for 'block' one can specify '{'block_size': (20, 20, 10) }'. """ + default_assignment_simplifications: bool = False + """ + If `True` default simplifications are first performed on the Assignments. If problems occur during the + simplification a warning will be thrown. + Furthermore, it is essential to know that this is a two-stage process. The first stage of the process acts + on the level of the `AssignmentCollection`. In this part, `create_simplification_strategy` + from pystencils.simplificationfactory will be used to apply optimisations like insertion of constants to + remove pressure from the registers. Thus the first part of the optimisations can only be executed if + an `AssignmentCollection` is passed. The second part of the optimisation acts on the level of each Assignment + individually. In this stage, all optimisations from `sympy.codegen.rewriting.optims_c99` are applied + to each Assignment. Thus this stage can also be applied if a list of Assignments is passed. + """ cpu_prepend_optimizations: List[Callable] = field(default_factory=list) """ List of extra optimizations to perform first on the AST. @@ -195,8 +209,8 @@ def create_domain_kernel(assignments: List[Assignment], *, config: CreateKernelC >>> import numpy as np >>> s, d = ps.fields('s, d: [2D]') >>> assignment = ps.Assignment(d[0,0], s[0, 1] + s[0, -1] + s[1, 0] + s[-1, 0]) - >>> config = ps.CreateKernelConfig(cpu_openmp=True) - >>> kernel_ast = ps.kernelcreation.create_domain_kernel([assignment], config=config) + >>> kernel_config = ps.CreateKernelConfig(cpu_openmp=True) + >>> kernel_ast = ps.kernelcreation.create_domain_kernel([assignment], config=kernel_config) >>> kernel = kernel_ast.compile() >>> d_arr = np.zeros([5, 5]) >>> kernel(d=d_arr, s=np.ones([5, 5])) @@ -207,6 +221,15 @@ def create_domain_kernel(assignments: List[Assignment], *, config: CreateKernelC [0., 4., 4., 4., 0.], [0., 0., 0., 0., 0.]]) """ + # --- applying first default simplifications + try: + if config.default_assignment_simplifications and isinstance(assignments, AssignmentCollection): + simplification = create_simplification_strategy() + assignments = simplification(assignments) + except Exception as e: + warnings.warn(f"It was not possible to apply the default pystencils optimisations to the " + f"AssignmentCollection due to the following problem :{e}") + # ---- Normalizing parameters split_groups = () if isinstance(assignments, AssignmentCollection): @@ -214,6 +237,13 @@ def create_domain_kernel(assignments: List[Assignment], *, config: CreateKernelC split_groups = assignments.simplification_hints['split_groups'] assignments = assignments.all_assignments + try: + if config.default_assignment_simplifications: + assignments = apply_sympy_optimisations(assignments) + except Exception as e: + warnings.warn(f"It was not possible to apply the default SymPy optimisations to the " + f"Assignments due to the following problem :{e}") + # ---- Creating ast ast = None if config.target == Target.CPU: @@ -304,9 +334,9 @@ def create_indexed_kernel(assignments: List[Assignment], *, config: CreateKernel >>> >>> # Additional values stored in index field can be accessed in the kernel as well >>> s, d = ps.fields('s, d: [2D]') - >>> assignment = ps.Assignment(d[0,0], 2 * s[0, 1] + 2 * s[1, 0] + idx_field('val')) - >>> config = ps.CreateKernelConfig(index_fields=[idx_field], coordinate_names=('x', 'y')) - >>> kernel_ast = ps.create_indexed_kernel([assignment], config=config) + >>> assignment = ps.Assignment(d[0, 0], 2 * s[0, 1] + 2 * s[1, 0] + idx_field('val')) + >>> kernel_config = ps.CreateKernelConfig(index_fields=[idx_field], coordinate_names=('x', 'y')) + >>> kernel_ast = ps.create_indexed_kernel([assignment], config=kernel_config) >>> kernel = kernel_ast.compile() >>> d_arr = np.zeros([5, 5]) >>> kernel(s=np.ones([5, 5]), d=d_arr, idx=index_arr) diff --git a/pystencils/simp/assignment_collection.py b/pystencils/simp/assignment_collection.py index 33102dee5085549f7ce827efe99d15f67da2008d..f493b08e931bf69f0b9713255c0eee87e9bace66 100644 --- a/pystencils/simp/assignment_collection.py +++ b/pystencils/simp/assignment_collection.py @@ -19,13 +19,14 @@ class AssignmentCollection: Additionally a dictionary of simplification hints is stored, which are set by the functions that create assignment collections to transport information to the simplification system. - Attributes: - main_assignments: list of assignments - subexpressions: list of assignments defining subexpressions used in main equations - simplification_hints: dict that is used to annotate the assignment collection with hints that are + Args: + main_assignments: List of assignments. Main assignments are characterised, that the right hand side of each + assignment is a field access. Thus the generated equations write on arrays. + subexpressions: List of assignments defining subexpressions used in main equations + simplification_hints: Dict that is used to annotate the assignment collection with hints that are used by the simplification system. See documentation of the simplification rules for potentially required hints and their meaning. - subexpression_symbol_generator: generator for new symbols that are used when new subexpressions are added + subexpression_symbol_generator: Generator for new symbols that are used when new subexpressions are added used to get new symbols that are unique for this AssignmentCollection """ @@ -33,9 +34,13 @@ class AssignmentCollection: # ------------------------------- Creation & Inplace Manipulation -------------------------------------------------- def __init__(self, main_assignments: Union[List[Assignment], Dict[sp.Expr, sp.Expr]], - subexpressions: Union[List[Assignment], Dict[sp.Expr, sp.Expr]] = {}, + subexpressions: Union[List[Assignment], Dict[sp.Expr, sp.Expr]] = None, simplification_hints: Optional[Dict[str, Any]] = None, subexpression_symbol_generator: Iterator[sp.Symbol] = None) -> None: + + if subexpressions is None: + subexpressions = {} + if isinstance(main_assignments, Dict): main_assignments = [Assignment(k, v) for k, v in main_assignments.items()] diff --git a/pystencils/simp/simplifications.py b/pystencils/simp/simplifications.py index 7345ad6f3fff3f79bc54237d0dcc76253831efa8..114b86a4013cccc3a74b641b0aabf747dbd3035d 100644 --- a/pystencils/simp/simplifications.py +++ b/pystencils/simp/simplifications.py @@ -3,9 +3,11 @@ from typing import Callable, List, Sequence, Union from collections import defaultdict import sympy as sp +from sympy.codegen.rewriting import optims_c99, optimize +from sympy.codegen.rewriting import ReplaceOptim from pystencils.assignment import Assignment -from pystencils.astnodes import Node +from pystencils.astnodes import Node, SympyAssignment from pystencils.field import AbstractField, Field from pystencils.sympyextensions import subs_additive, is_constant, recursive_collect @@ -223,3 +225,24 @@ def apply_on_all_subexpressions(operation: Callable[[sp.Expr], sp.Expr]): f.__name__ = operation.__name__ return f + + +def apply_sympy_optimisations(assignments): + """ Evaluates constant expressions (e.g. :math:`\\sqrt{3}` will be replaced by its floating point representation) + and applies the default sympy optimisations. See sympy.codegen.rewriting + """ + + # 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()) + + sympy_optimisations = [evaluate_constant_terms] + list(optims_c99) + + assignments = [Assignment(a.lhs, optimize(a.rhs, sympy_optimisations)) + if hasattr(a, 'lhs') + else a for a in assignments] + assignments_nodes = [a.atoms(SympyAssignment) for a in assignments] + for a in chain.from_iterable(assignments_nodes): + a.optimize(sympy_optimisations) + + return assignments diff --git a/pystencils/simp/subexpression_insertion.py b/pystencils/simp/subexpression_insertion.py index 9293b56be8f6fb223da6778a33db7e07bf463fc4..27ac798269bb043c5c8361ea019492a2be5dd9a3 100644 --- a/pystencils/simp/subexpression_insertion.py +++ b/pystencils/simp/subexpression_insertion.py @@ -4,7 +4,7 @@ from pystencils.sympyextensions import is_constant # Subexpression Insertion -def insert_subexpressions(ac, selection_callback, skip=set()): +def insert_subexpressions(ac, selection_callback, skip=None): """ Removes a number of subexpressions from an assignment collection by inserting their right-hand side wherever they occur. @@ -16,6 +16,8 @@ def insert_subexpressions(ac, selection_callback, skip=set()): - skip: Set of symbols (left-hand sides of subexpressions) that should be ignored even if qualified by the callback. """ + if skip is None: + skip = set() i = 0 while i < len(ac.subexpressions): exp = ac.subexpressions[i] diff --git a/pystencils/simplificationfactory.py b/pystencils/simplificationfactory.py new file mode 100644 index 0000000000000000000000000000000000000000..50ee2d7f8175ff7445b9ec0f13071164c50099b1 --- /dev/null +++ b/pystencils/simplificationfactory.py @@ -0,0 +1,18 @@ +from pystencils.simp import (SimplificationStrategy, insert_constants, insert_symbol_times_minus_one, + insert_constant_multiples, insert_constant_additions, insert_squares, insert_zeros) + + +def create_simplification_strategy(): + """ + Creates a default simplification `ps.simp.SimplificationStrategy`. The idea behind the default simplification + strategy is to reduce the number of subexpressions by inserting single constants and to evaluate constant + terms beforehand. + """ + s = SimplificationStrategy() + s.add(insert_symbol_times_minus_one) + s.add(insert_constant_multiples) + s.add(insert_constant_additions) + s.add(insert_squares) + s.add(insert_zeros) + s.add(insert_constants) + s.add(lambda ac: ac.new_without_unused_subexpressions()) diff --git a/pystencils/sympyextensions.py b/pystencils/sympyextensions.py index 29b524eef506dcabcdac3c487bbf53e87965b42a..f63328d81781c578786c0019edb1549b83a53c83 100644 --- a/pystencils/sympyextensions.py +++ b/pystencils/sympyextensions.py @@ -453,7 +453,7 @@ def recursive_collect(expr, symbols, order_by_occurences=False): return rec_sum -def count_operations(term: Union[sp.Expr, List[sp.Expr]], +def count_operations(term: Union[sp.Expr, List[sp.Expr], List[Assignment]], only_type: Optional[str] = 'real') -> Dict[str, int]: """Counts the number of additions, multiplications and division. diff --git a/pystencils_tests/test_Min_Max.py b/pystencils_tests/test_Min_Max.py index 18cd2d99d10ec0b03a38630d07679f3191281493..c227fbf149bac148ff4a497f77e28a9e33d5aada 100644 --- a/pystencils_tests/test_Min_Max.py +++ b/pystencils_tests/test_Min_Max.py @@ -1,10 +1,13 @@ -import sympy +import pytest + +import sympy as sp import numpy import pystencils from pystencils.datahandling import create_data_handling -def test_max(): +@pytest.mark.parametrize('sympy_function', [sp.Min, sp.Max]) +def test_max(sympy_function): dh = create_data_handling(domain_size=(10, 10), periodicity=True) x = dh.add_array('x', values_per_cell=1) @@ -15,56 +18,28 @@ def test_max(): dh.fill("z", 2.0, ghost_layers=True) # test sp.Max with one argument - assignment_1 = pystencils.Assignment(x.center, sympy.Max(y.center + 3.3)) + assignment_1 = pystencils.Assignment(x.center, sympy_function(y.center + 3.3)) ast_1 = pystencils.create_kernel(assignment_1) kernel_1 = ast_1.compile() # test sp.Max with two arguments - assignment_2 = pystencils.Assignment(x.center, sympy.Max(0.5, y.center - 1.5)) + assignment_2 = pystencils.Assignment(x.center, sympy_function(0.5, y.center - 1.5)) ast_2 = pystencils.create_kernel(assignment_2) kernel_2 = ast_2.compile() # test sp.Max with many arguments - assignment_3 = pystencils.Assignment(x.center, sympy.Max(z.center, 4.5, y.center - 1.5, y.center + z.center)) + assignment_3 = pystencils.Assignment(x.center, sympy_function(z.center, 4.5, y.center - 1.5, y.center + z.center)) ast_3 = pystencils.create_kernel(assignment_3) kernel_3 = ast_3.compile() - dh.run_kernel(kernel_1) - assert numpy.all(dh.cpu_arrays["x"] == 4.3) - dh.run_kernel(kernel_2) - assert numpy.all(dh.cpu_arrays["x"] == 0.5) - dh.run_kernel(kernel_3) - assert numpy.all(dh.cpu_arrays["x"] == 4.5) - - -def test_min(): - dh = create_data_handling(domain_size=(10, 10), periodicity=True) - - x = dh.add_array('x', values_per_cell=1) - dh.fill("x", 0.0, ghost_layers=True) - y = dh.add_array('y', values_per_cell=1) - dh.fill("y", 1.0, ghost_layers=True) - z = dh.add_array('z', values_per_cell=1) - dh.fill("z", 2.0, ghost_layers=True) - - # test sp.Min with one argument - assignment_1 = pystencils.Assignment(x.center, sympy.Min(y.center + 3.3)) - ast_1 = pystencils.create_kernel(assignment_1) - kernel_1 = ast_1.compile() - - # test sp.Min with two arguments - assignment_2 = pystencils.Assignment(x.center, sympy.Min(0.5, y.center - 1.5)) - ast_2 = pystencils.create_kernel(assignment_2) - kernel_2 = ast_2.compile() - - # test sp.Min with many arguments - assignment_3 = pystencils.Assignment(x.center, sympy.Min(z.center, 4.5, y.center - 1.5, y.center + z.center)) - ast_3 = pystencils.create_kernel(assignment_3) - kernel_3 = ast_3.compile() + if sympy_function is sp.Max: + results = [4.3, 0.5, 4.5] + else: + results = [4.3, -0.5, -0.5] dh.run_kernel(kernel_1) - assert numpy.all(dh.cpu_arrays["x"] == 4.3) + assert numpy.all(dh.gather_array('x') == results[0]) dh.run_kernel(kernel_2) - assert numpy.all(dh.cpu_arrays["x"] == - 0.5) + assert numpy.all(dh.gather_array('x') == results[1]) dh.run_kernel(kernel_3) - assert numpy.all(dh.cpu_arrays["x"] == - 0.5) + assert numpy.all(dh.gather_array('x') == results[2]) diff --git a/pystencils_tests/test_astnodes.py b/pystencils_tests/test_astnodes.py index 385c4e2234f7879d61ef64cb7dbe4961071e09c6..688f63ed95e3d7f40feb428f8ca778e0e38d9288 100644 --- a/pystencils_tests/test_astnodes.py +++ b/pystencils_tests/test_astnodes.py @@ -1,24 +1,19 @@ import pytest +import sys import sympy as sp import pystencils as ps from pystencils import Assignment from pystencils.astnodes import Block, LoopOverCoordinate, SkipIteration, SympyAssignment -sympy_numeric_version = [int(x, 10) for x in sp.__version__.split('.') if x.isdigit()] -if len(sympy_numeric_version) < 3: - sympy_numeric_version.append(0) -sympy_numeric_version.reverse() -sympy_version = sum(x * (100 ** i) for i, x in enumerate(sympy_numeric_version)) - dst = ps.fields('dst(8): double[2D]') s = sp.symbols('s_:8') x = sp.symbols('x') y = sp.symbols('y') +python_version = f"{sys.version_info.major}.{sys.version_info.minor}.{sys.version_info.micro}" + -@pytest.mark.skipif(sympy_version < 10501, - reason="Old Sympy Versions behave differently which wont be supported in the near future") def test_kernel_function(): assignments = [ Assignment(dst[0, 0](0), s[0]), @@ -44,8 +39,6 @@ def test_skip_iteration(): assert skipped.undefined_symbols == set() -@pytest.mark.skipif(sympy_version < 10501, - reason="Old Sympy Versions behave differently which wont be supported in the near future") def test_block(): assignments = [ Assignment(dst[0, 0](0), s[0]), @@ -92,17 +85,23 @@ def test_loop_over_coordinate(): assert loop.step == 2 -def test_sympy_assignment(): - pytest.importorskip('sympy.codegen.rewriting') - from sympy.codegen.rewriting import optims_c99 +@pytest.mark.parametrize('default_assignment_simplifications', [False, True]) +@pytest.mark.skipif(python_version == '3.8.2', reason="For this python version a strange bug in mpmath occurs") +def test_sympy_assignment(default_assignment_simplifications): assignment = SympyAssignment(dst[0, 0](0), sp.log(x + 3) / sp.log(2) + sp.log(x ** 2 + 1)) - assignment.optimize(optims_c99) - ast = ps.create_kernel([assignment]) + config = ps.CreateKernelConfig(default_assignment_simplifications=default_assignment_simplifications) + ast = ps.create_kernel([assignment], config=config) code = ps.get_code_str(ast) - - assert 'log1p' in code - assert 'log2' in code + + if default_assignment_simplifications: + assert 'log1p' in code + # constant term is directly evaluated + assert 'log2' not in code + else: + # no optimisations will be applied so the optimised version of log will not be in the code + assert 'log1p' not in code + assert 'log2' not in code assignment.replace(assignment.lhs, dst[0, 0](1)) assignment.replace(assignment.rhs, sp.log(2)) diff --git a/pystencils_tests/test_simplifications.py b/pystencils_tests/test_simplifications.py index 62fa262fb203d02b2c10e624893b428f97d33328..1c9ed3c0cc88e34fed30fe0f115c4e8afa3dc281 100644 --- a/pystencils_tests/test_simplifications.py +++ b/pystencils_tests/test_simplifications.py @@ -1,5 +1,7 @@ +from sys import version_info as vs import pytest import sympy as sp +import pystencils as ps from pystencils.simp import subexpression_substitution_in_main_assignments from pystencils.simp import add_subexpressions_for_divisions @@ -136,3 +138,52 @@ def test_add_subexpressions_for_field_reads(): assert len(ac.subexpressions) == 0 ac = add_subexpressions_for_field_reads(ac) assert len(ac.subexpressions) == 2 + + +@pytest.mark.parametrize('target', (ps.Target.CPU, ps.Target.GPU)) +@pytest.mark.parametrize('simplification', (True, False)) +@pytest.mark.skipif((vs.major, vs.minor, vs.micro) == (3, 8, 2), reason="does not work on python 3.8.2 for some reason") +def test_sympy_optimizations(target, simplification): + if target == ps.Target.GPU: + pytest.importorskip("pycuda") + src, dst = ps.fields('src, dst: float32[2d]') + + # Triggers Sympy's expm1 optimization + # Sympy's expm1 optimization is tedious to use and the behaviour is highly depended on the sympy version. In + # some cases the exp expression has to be encapsulated in brackets or multiplied with 1 or 1.0 + # for sympy to work properly ... + assignments = ps.AssignmentCollection({ + src[0, 0]: 1.0 * (sp.exp(dst[0, 0]) - 1) + }) + + config = ps.CreateKernelConfig(target=target, default_assignment_simplifications=simplification) + ast = ps.create_kernel(assignments, config=config) + + code = ps.get_code_str(ast) + if simplification: + assert 'expm1(' in code + else: + assert 'expm1(' not in code + + +@pytest.mark.parametrize('target', (ps.Target.CPU, ps.Target.GPU)) +@pytest.mark.parametrize('simplification', (True, False)) +@pytest.mark.skipif((vs.major, vs.minor, vs.micro) == (3, 8, 2), reason="does not work on python 3.8.2 for some reason") +def test_evaluate_constant_terms(target, simplification): + if target == ps.Target.GPU: + pytest.importorskip("pycuda") + src, dst = ps.fields('src, dst: float32[2d]') + + # Triggers Sympy's cos optimization + assignments = ps.AssignmentCollection({ + src[0, 0]: -sp.cos(1) + dst[0, 0] + }) + + config = ps.CreateKernelConfig(target=target, default_assignment_simplifications=simplification) + ast = ps.create_kernel(assignments, config=config) + code = ps.get_code_str(ast) + if simplification: + assert 'cos(' not in code + else: + assert 'cos(' in code + print(code) diff --git a/pystencils_tests/test_sum_prod.py b/pystencils_tests/test_sum_prod.py index 4b4cd7131ad4b860c754865cd35f3960399e9b5a..2f6bf7359ad9d74c85b5d679dab772b5d7c00803 100644 --- a/pystencils_tests/test_sum_prod.py +++ b/pystencils_tests/test_sum_prod.py @@ -7,34 +7,36 @@ """ """ +import pytest import numpy as np -import sympy -from sympy.abc import k +import sympy as sp +import sympy.abc -import pystencils +import pystencils as ps from pystencils.data_types import create_type -def test_sum(): +@pytest.mark.parametrize('default_assignment_simplifications', [False, True]) +def test_sum(default_assignment_simplifications): - sum = sympy.Sum(k, (k, 1, 100)) + sum = sp.Sum(sp.abc.k, (sp.abc.k, 1, 100)) expanded_sum = sum.doit() print(sum) print(expanded_sum) - x = pystencils.fields('x: float32[1d]') + x = ps.fields('x: float32[1d]') - assignments = pystencils.AssignmentCollection({ - x.center(): sum - }) + assignments = ps.AssignmentCollection({x.center(): sum}) - ast = pystencils.create_kernel(assignments) - code = str(pystencils.get_code_obj(ast)) + config = ps.CreateKernelConfig(default_assignment_simplifications=default_assignment_simplifications) + ast = ps.create_kernel(assignments, config=config) + code = ps.get_code_str(ast) kernel = ast.compile() print(code) - assert 'double sum' in code + if default_assignment_simplifications is False: + assert 'double sum' in code array = np.zeros((10,), np.float32) @@ -43,27 +45,28 @@ def test_sum(): assert np.allclose(array, int(expanded_sum) * np.ones_like(array)) -def test_sum_use_float(): +@pytest.mark.parametrize('default_assignment_simplifications', [False, True]) +def test_sum_use_float(default_assignment_simplifications): - sum = sympy.Sum(k, (k, 1, 100)) + sum = sympy.Sum(sp.abc.k, (sp.abc.k, 1, 100)) expanded_sum = sum.doit() print(sum) print(expanded_sum) - x = pystencils.fields('x: float32[1d]') + x = ps.fields('x: float32[1d]') - assignments = pystencils.AssignmentCollection({ - x.center(): sum - }) + assignments = ps.AssignmentCollection({x.center(): sum}) - ast = pystencils.create_kernel(assignments, data_type=create_type('float32')) - code = str(pystencils.get_code_obj(ast)) + config = ps.CreateKernelConfig(default_assignment_simplifications=default_assignment_simplifications, + data_type=create_type('float32')) + ast = ps.create_kernel(assignments, config=config) + code = ps.get_code_str(ast) kernel = ast.compile() print(code) - print(pystencils.get_code_obj(ast)) - assert 'float sum' in code + if default_assignment_simplifications is False: + assert 'float sum' in code array = np.zeros((10,), np.float32) @@ -72,9 +75,10 @@ def test_sum_use_float(): assert np.allclose(array, int(expanded_sum) * np.ones_like(array)) -def test_product(): +@pytest.mark.parametrize('default_assignment_simplifications', [False, True]) +def test_product(default_assignment_simplifications): - k = pystencils.TypedSymbol('k', create_type('int64')) + k = ps.TypedSymbol('k', create_type('int64')) sum = sympy.Product(k, (k, 1, 10)) expanded_sum = sum.doit() @@ -82,18 +86,19 @@ def test_product(): print(sum) print(expanded_sum) - x = pystencils.fields('x: int64[1d]') + x = ps.fields('x: int64[1d]') - assignments = pystencils.AssignmentCollection({ - x.center(): sum - }) + assignments = ps.AssignmentCollection({x.center(): sum}) - ast = pystencils.create_kernel(assignments) - code = pystencils.get_code_str(ast) + config = ps.CreateKernelConfig(default_assignment_simplifications=default_assignment_simplifications) + + ast = ps.create_kernel(assignments, config=config) + code = ps.get_code_str(ast) kernel = ast.compile() print(code) - assert 'int64_t product' in code + if default_assignment_simplifications is False: + assert 'int64_t product' in code array = np.zeros((10,), np.int64) @@ -104,8 +109,8 @@ def test_product(): def test_prod_var_limit(): - k = pystencils.TypedSymbol('k', create_type('int64')) - limit = pystencils.TypedSymbol('limit', create_type('int64')) + k = ps.TypedSymbol('k', create_type('int64')) + limit = ps.TypedSymbol('limit', create_type('int64')) sum = sympy.Sum(k, (k, 1, limit)) expanded_sum = sum.replace(limit, 100).doit() @@ -113,14 +118,12 @@ def test_prod_var_limit(): print(sum) print(expanded_sum) - x = pystencils.fields('x: int64[1d]') + x = ps.fields('x: int64[1d]') - assignments = pystencils.AssignmentCollection({ - x.center(): sum - }) + assignments = ps.AssignmentCollection({x.center(): sum}) - ast = pystencils.create_kernel(assignments) - pystencils.show_code(ast) + ast = ps.create_kernel(assignments) + ps.show_code(ast) kernel = ast.compile() array = np.zeros((10,), np.int64)