diff --git a/pystencils/kernel_contrains_check.py b/pystencils/kernel_contrains_check.py index a39c7d0ca12d06acfcf413821b467e9a4270b5a5..7f7756d6f0318930c9e4d85656ce6314146da666 100644 --- a/pystencils/kernel_contrains_check.py +++ b/pystencils/kernel_contrains_check.py @@ -58,6 +58,7 @@ class KernelConstraintsCheck: if obj.false_block: self.visit(obj.false_block) self.process_expression(obj.condition_expr) + self.process_expression(obj.true_block) self.check_double_write_condition = old_double_write self.scopes.pop() elif isinstance(obj, ast.Block): diff --git a/pystencils/simp/simplifications.py b/pystencils/simp/simplifications.py index 720abb52ad0f66e030fa7b5f922b8ab0771124bd..086ff984115d5fcf03ed49dac45850cb17bd8c62 100644 --- a/pystencils/simp/simplifications.py +++ b/pystencils/simp/simplifications.py @@ -3,12 +3,10 @@ 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, SympyAssignment -from pystencils.field import AbstractField, Field +from pystencils.astnodes import Node +from pystencils.field import Field from pystencils.sympyextensions import subs_additive, is_constant, recursive_collect @@ -164,7 +162,7 @@ def add_subexpressions_for_sums(ac): for eq in ac.all_assignments: search_addends(eq.rhs) - addends = [a for a in addends if not isinstance(a, sp.Symbol) or isinstance(a, AbstractField.AbstractAccess)] + addends = [a for a in addends if not isinstance(a, sp.Symbol) or isinstance(a, Field.Access)] new_symbol_gen = ac.subexpression_symbol_generator substitutions = {addend: new_symbol for new_symbol, addend in zip(new_symbol_gen, addends)} return ac.new_with_substitutions(substitutions, True, substitute_on_lhs=False) @@ -226,23 +224,29 @@ 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(17)) - - 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 +# TODO Markus +# TODO: make this really work for Assignmentcollections +# TODO: this function should ONLY evaluate +# TODO: do the optims_c99 elsewhere optionally +# def apply_sympy_optimisations(ac: AssignmentCollection): +# """ 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 +# +# assignments = ac.all_assignments +# +# 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 AssignmentCollection(assignments) diff --git a/pystencils/typing/leaf_typing.py b/pystencils/typing/leaf_typing.py index d8f41fd1847fda46157dc7191864021c3adb4d84..f92a0db73df6f787a715b35478c7f406d10e78b5 100644 --- a/pystencils/typing/leaf_typing.py +++ b/pystencils/typing/leaf_typing.py @@ -250,7 +250,7 @@ class TypeAdder: type(rhs) in pystencils.integer_functions.__dict__.values(): new_args = [self.process_expression(a, type_constants) for a in rhs.args] types_of_expressions = [get_type_of_expression(a) for a in new_args] - arg_type = collate_types(types_of_expressions, forbid_collation_to_float=True) + arg_type = collate_types(types_of_expressions) new_args = [a if not hasattr(a, 'dtype') or a.dtype == arg_type else CastFunc(a, arg_type) for a in new_args] diff --git a/pystencils_tests/test_fvm.py b/pystencils_tests/test_fvm.py index f4c0a663ea069bcbbb2d694e82c14259adc97bfb..55ce1f7db2e859978a4afb5bb3a763ca5e7d2d56 100644 --- a/pystencils_tests/test_fvm.py +++ b/pystencils_tests/test_fvm.py @@ -240,10 +240,12 @@ def advection_diffusion_fluctuations(dim: int): advection_diffusion_fluctuations.runners = {} -@pytest.mark.parametrize("velocity", list(product([0, 0.00041], [0, -0.00031]))) -@pytest.mark.parametrize("density", [27.0, 56.5]) -@pytest.mark.longrun -def test_advection_diffusion_fluctuation_2d(density, velocity): +# @pytest.mark.parametrize("velocity", list(product([0, 0.00041], [0, -0.00031]))) +# @pytest.mark.parametrize("density", [27.0, 56.5]) +# @pytest.mark.longrun +def test_advection_diffusion_fluctuation_2d(): + density = 27.0 + velocity = [0, 0.00041] if 2 not in advection_diffusion_fluctuations.runners: advection_diffusion_fluctuations.runners[2] = advection_diffusion_fluctuations(2) advection_diffusion_fluctuations.runners[2](density, velocity) diff --git a/pystencils_tests/test_random.py b/pystencils_tests/test_random.py index b29f15eb7d94a155b93f9191f2d75b744c262293..2310f4c1dd06e33b1c2a9d0aa34e281f6a804252 100644 --- a/pystencils_tests/test_random.py +++ b/pystencils_tests/test_random.py @@ -3,6 +3,7 @@ import numpy as np import pytest import pystencils as ps +from pystencils.astnodes import SympyAssignment from pystencils.rng import PhiloxFourFloats, PhiloxTwoDoubles, AESNIFourFloats, AESNITwoDoubles, random_symbol from pystencils.backends.simd_instruction_sets import get_supported_instruction_sets from pystencils.cpu.cpujit import get_compiler_config @@ -22,8 +23,7 @@ if get_compiler_config()['os'] == 'windows': instruction_sets.remove('avx512') -@pytest.mark.parametrize('target,rng', ( -(Target.CPU, 'philox'), (Target.CPU, 'aesni'), (Target.GPU, 'philox'))) +@pytest.mark.parametrize('target,rng', ((Target.CPU, 'philox'), (Target.CPU, 'aesni'), (Target.GPU, 'philox'))) @pytest.mark.parametrize('precision', ('float', 'double')) @pytest.mark.parametrize('dtype', ('float', 'double')) def test_rng(target, rng, precision, dtype, t=124, offsets=(0, 0), keys=(0, 0), offset_values=None): @@ -42,7 +42,7 @@ def test_rng(target, rng, precision, dtype, t=124, offsets=(0, 0), keys=(0, 0), dh.fill(f.name, 42.0) rng_node = RNGs[(rng, precision)](dh.dim, offsets=offsets, keys=keys) - assignments = [rng_node] + [ps.Assignment(f(i), s) for i, s in enumerate(rng_node.result_symbols)] + assignments = [rng_node] + [SympyAssignment(f(i), s) for i, s in enumerate(rng_node.result_symbols)] kernel = ps.create_kernel(assignments, target=dh.default_target).compile() dh.all_to_gpu() @@ -130,7 +130,7 @@ def test_rng_vectorized(target, rng, precision, dtype, t=130, offsets=(1, 3), ke ref = dh.add_array("ref", values_per_cell=4 if precision == 'float' else 2) rng_node = RNGs[(rng, precision)](dh.dim, offsets=offsets) - assignments = [rng_node] + [ps.Assignment(ref(i), s) for i, s in enumerate(rng_node.result_symbols)] + assignments = [rng_node] + [SympyAssignment(ref(i), s) for i, s in enumerate(rng_node.result_symbols)] kernel = ps.create_kernel(assignments, target=dh.default_target).compile() kwargs = {'time_step': t} @@ -139,7 +139,7 @@ def test_rng_vectorized(target, rng, precision, dtype, t=130, offsets=(1, 3), ke dh.run_kernel(kernel, **kwargs) rng_node = RNGs[(rng, precision)](dh.dim, offsets=offsets) - assignments = [rng_node] + [ps.Assignment(f(i), s) for i, s in enumerate(rng_node.result_symbols)] + assignments = [rng_node] + [SympyAssignment(f(i), s) for i, s in enumerate(rng_node.result_symbols)] kernel = ps.create_kernel(assignments, target=dh.default_target, cpu_vectorize_info=cpu_vectorize_info).compile() dh.run_kernel(kernel, **kwargs) @@ -153,14 +153,13 @@ def test_rng_vectorized(target, rng, precision, dtype, t=130, offsets=(1, 3), ke @pytest.mark.parametrize('vectorized', (False, True)) def test_rng_symbol(vectorized): """Make sure that the RNG symbol generator generates symbols and that the resulting code compiles""" + cpu_vectorize_info = None if vectorized: if not instruction_sets: pytest.skip("cannot detect CPU instruction set") else: cpu_vectorize_info = {'assume_inner_stride_one': True, 'assume_aligned': True, 'instruction_set': instruction_sets[-1]} - else: - cpu_vectorize_info = None dh = ps.create_data_handling((8, 8), default_ghost_layers=0, default_target=Target.CPU) f = dh.add_array("f", values_per_cell=2 * dh.dim, alignment=True) diff --git a/pystencils_tests/test_types.py b/pystencils_tests/test_types.py index b6a7cd81cf8b7618ab69f6e0dd69094f93de3238..8ac96f84e732debd8e1ff50426c483b6b5523868 100644 --- a/pystencils_tests/test_types.py +++ b/pystencils_tests/test_types.py @@ -1,24 +1,93 @@ +import pytest + +import pystencils.config import sympy as sp import numpy as np import pystencils as ps -from pystencils import data_types -from pystencils.data_types import TypedSymbol, get_type_of_expression, VectorType, collate_types, create_type, \ - typed_symbols, type_all_numbers, matrix_symbols, cast_func, pointer_arithmetic_func, PointerType +from pystencils.typing import TypedSymbol, get_type_of_expression, VectorType, collate_types, \ + typed_symbols, CastFunc, PointerArithmeticFunc, PointerType, result_type, BasicType + + +def test_result_type(): + i = np.dtype('int32') + l = np.dtype('int64') + ui = np.dtype('uint32') + ul = np.dtype('uint64') + f = np.dtype('float32') + d = np.dtype('float64') + b = np.dtype('bool') + + assert result_type(i, l) == l + assert result_type(l, i) == l + assert result_type(ui, i) == i + assert result_type(ui, l) == l + assert result_type(ul, i) == i + assert result_type(ul, l) == l + assert result_type(d, f) == d + assert result_type(f, d) == d + assert result_type(i, f) == f + assert result_type(l, f) == f + assert result_type(ui, f) == f + assert result_type(ul, f) == f + assert result_type(i, d) == d + assert result_type(l, d) == d + assert result_type(ui, d) == d + assert result_type(ul, d) == d + assert result_type(b, i) == i + assert result_type(b, l) == l + assert result_type(b, ui) == ui + assert result_type(b, ul) == ul + assert result_type(b, f) == f + assert result_type(b, d) == d + + +@pytest.mark.parametrize('dtype', ('float64', 'float32', 'int64', 'int32', 'uint32', 'uint64')) +def test_simple_add(dtype): + constant = 1.0 + if dtype[0] in 'ui': + constant = 1 + f = ps.fields(f"f: {dtype}[1D]") + d = TypedSymbol("d", dtype) + + test_arr = np.array([constant], dtype=dtype) + + ur = ps.Assignment(f[0], f[0] + d) + + ast = ps.create_kernel(ur) + code = ps.get_code_str(ast) + kernel = ast.compile() + kernel(f=test_arr, d=constant) + + assert test_arr[0] == constant+constant + + +@pytest.mark.parametrize('dtype1', ('float64', 'float32', 'int64', 'int32', 'uint32', 'uint64')) +@pytest.mark.parametrize('dtype2', ('float64', 'float32', 'int64', 'int32', 'uint32', 'uint64')) +def test_mixed_add(dtype1, dtype2): + + constant = 1 + f = ps.fields(f"f: {dtype1}[1D]") + g = ps.fields(f"g: {dtype2}[1D]") + test_f = np.array([constant], dtype=dtype1) + test_g = np.array([constant], dtype=dtype2) -def test_parsing(): - assert str(data_types.create_composite_type_from_string("const double *")) == "double const *" - assert str(data_types.create_composite_type_from_string("double const *")) == "double const *" + ur = ps.Assignment(f[0], f[0] + g[0]) - t1 = data_types.create_composite_type_from_string("const double * const * const restrict") - t2 = data_types.create_composite_type_from_string(str(t1)) - assert t1 == t2 + # TODO Markus: check for the logging if colate_types(dtype1, dtype2) != dtype1 + ast = ps.create_kernel(ur) + code = ps.get_code_str(ast) + kernel = ast.compile() + kernel(f=test_f, g=test_g) + assert test_f[0] == constant+constant + +# TODO vector def test_collation(): - double_type = create_type("double") - float_type = create_type("float32") + double_type = BasicType('float64') + float_type = BasicType('float32') double4_type = VectorType(double_type, 4) float4_type = VectorType(float_type, 4) assert collate_types([double_type, float_type]) == double_type @@ -27,20 +96,23 @@ def test_collation(): def test_vector_type(): - double_type = create_type("double") - float_type = create_type("float32") + double_type = BasicType('float64') + float_type = BasicType('float32') double4_type = VectorType(double_type, 4) float4_type = VectorType(float_type, 4) assert double4_type.item_size == 4 assert float4_type.item_size == 4 - assert not double4_type == 4 + double4_type2 = VectorType(double_type, 4) + assert double4_type == double4_type2 + assert double4_type != 4 + assert double4_type != float4_type def test_pointer_type(): - double_type = create_type("double") - float_type = create_type("float32") + double_type = BasicType('float64') + float_type = BasicType('float32') double4_type = PointerType(double_type, restrict=True) float4_type = PointerType(float_type, restrict=False) @@ -72,96 +144,104 @@ def test_assumptions(): assert x.shape[0].is_nonnegative assert (2 * x.shape[0]).is_nonnegative assert (2 * x.shape[0]).is_integer - assert (TypedSymbol('a', create_type('uint64'))).is_nonnegative - assert (TypedSymbol('a', create_type('uint64'))).is_positive is None - assert (TypedSymbol('a', create_type('uint64')) + 1).is_positive + assert (TypedSymbol('a', BasicType('uint64'))).is_nonnegative + assert (TypedSymbol('a', BasicType('uint64'))).is_positive is None + assert (TypedSymbol('a', BasicType('uint64')) + 1).is_positive assert (x.shape[0] + 1).is_real -def test_sqrt_of_integer(): +@pytest.mark.parametrize('dtype', ('float64', 'float32')) +def test_sqrt_of_integer(dtype): """Regression test for bug where sqrt(3) was classified as integer""" - f = ps.fields("f: [1D]") - tmp = sp.symbols("tmp") + f = ps.fields(f'f: {dtype}[1D]') + tmp = sp.symbols('tmp') assignments = [ps.Assignment(tmp, sp.sqrt(3)), ps.Assignment(f[0], tmp)] - arr_double = np.array([1], dtype=np.float64) - kernel = ps.create_kernel(assignments).compile() - kernel(f=arr_double) - assert 1.7 < arr_double[0] < 1.8 + arr = np.array([1], dtype=dtype) + # TODO Jupyter add auto lhs float/double problem + config = pystencils.config.CreateKernelConfig(data_type=dtype, default_number_float=dtype) - f = ps.fields("f: float32[1D]") - tmp = sp.symbols("tmp") + ast = ps.create_kernel(assignments, config=config) + kernel = ast.compile() + kernel(f=arr) + assert 1.7 < arr[0] < 1.8 - assignments = [ps.Assignment(tmp, sp.sqrt(3)), - ps.Assignment(f[0], tmp)] - arr_single = np.array([1], dtype=np.float32) - config = ps.CreateKernelConfig(data_type="float32") - kernel = ps.create_kernel(assignments, config=config).compile() - kernel(f=arr_single) - - code = ps.get_code_str(kernel.ast) - # ps.show_code(kernel.ast) - # 1.7320508075688772935 --> it is actually correct to round to ...773. This was wrong before !282 - assert "1.7320508075688773f" in code - assert 1.7 < arr_single[0] < 1.8 + code = ps.get_code_str(ast) + constant = '1.7320508075688772f' + if dtype == 'float32': + assert constant in code + else: + assert constant not in code -def test_integer_comparision(): - f = ps.fields("f [2D]") - d = sp.Symbol("dir") +@pytest.mark.parametrize('dtype', ('float64', 'float32')) +def test_integer_comparision(dtype): + f = ps.fields(f"f: {dtype}[2D]") + d = TypedSymbol("dir", "int64") ur = ps.Assignment(f[0, 0], sp.Piecewise((0, sp.Equality(d, 1)), (f[0, 0], True))) ast = ps.create_kernel(ur) code = ps.get_code_str(ast) - assert "_data_f_00[_stride_f_1*ctr_1] = ((((dir) == (1))) ? (0.0): (_data_f_00[_stride_f_1*ctr_1]));" in code + # There should be an explicit cast for the integer zero to the type of the field on the rhs + if dtype == 'float64': + t = "_data_f_00[_stride_f_1*ctr_1] = ((((dir) == (1))) ? (0.0): (_data_f_00[_stride_f_1*ctr_1]));" + else: + t = "_data_f_00[_stride_f_1*ctr_1] = ((((dir) == (1))) ? (0.0f): (_data_f_00[_stride_f_1*ctr_1]));" + assert t in code -def test_Basic_data_type(): +def test_typed_symbols_dtype(): assert typed_symbols(("s", "f"), np.uint) == typed_symbols("s, f", np.uint) t_symbols = typed_symbols(("s", "f"), np.uint) s = t_symbols[0] assert t_symbols[0] == TypedSymbol("s", np.uint) assert s.dtype.is_uint() - assert s.dtype.is_complex() == 0 - assert typed_symbols("s", str).dtype.is_other() - assert typed_symbols("s", bool).dtype.is_other() - assert typed_symbols("s", np.void).dtype.is_other() - - assert typed_symbols("s", np.float64).dtype.base_name == 'double' - # removed for old sympy version - # assert typed_symbols(("s"), np.float64).dtype.sympy_dtype == typed_symbols(("s"), float).dtype.sympy_dtype - - f, g = ps.fields("f, g : double[2D]") - - expr = ps.Assignment(f.center(), 2 * g.center() + 5) - new_expr = type_all_numbers(expr, np.float64) - - assert "cast_func(2, double)" in str(new_expr) - assert "cast_func(5, double)" in str(new_expr) - - m = matrix_symbols("a, b", np.uint, 3, 3) - assert len(m) == 2 - m = m[0] - for i, elem in enumerate(m): - assert elem == TypedSymbol(f"a{i}", np.uint) - assert elem.dtype.is_uint() + assert typed_symbols("s", np.float64).dtype.c_name == 'double' + assert typed_symbols("s", np.float32).dtype.c_name == 'float' assert TypedSymbol("s", np.uint).canonical == TypedSymbol("s", np.uint) assert TypedSymbol("s", np.uint).reversed == TypedSymbol("s", np.uint) def test_cast_func(): - assert cast_func(TypedSymbol("s", np.uint), np.int64).canonical == TypedSymbol("s", np.uint).canonical + assert CastFunc(TypedSymbol("s", np.uint), np.int64).canonical == TypedSymbol("s", np.uint).canonical - a = cast_func(5, np.uint) + a = CastFunc(5, np.uint) assert a.is_negative is False assert a.is_nonnegative def test_pointer_arithmetic_func(): - assert pointer_arithmetic_func(TypedSymbol("s", np.uint), 1).canonical == TypedSymbol("s", np.uint).canonical + assert PointerArithmeticFunc(TypedSymbol("s", np.uint), 1).canonical == TypedSymbol("s", np.uint).canonical + + +def test_division(): + f = ps.fields('f(10): float32[2D]') + m, tau = sp.symbols("m, tau") + + up = [ps.Assignment(tau, 1 / (0.5 + (3.0 * m))), + ps.Assignment(f.center, tau)] + config = pystencils.config.CreateKernelConfig(data_type='float32', default_number_float='float32') + ast = ps.create_kernel(up, config=config) + code = ps.get_code_str(ast) + + assert "((1.0f) / (m*3.0f + 0.5f))" in code + + +def test_pow(): + f = ps.fields('f(10): float32[2D]') + m, tau = sp.symbols("m, tau") + + up = [ps.Assignment(tau, m ** 1.5), + ps.Assignment(f.center, tau)] + + config = pystencils.config.CreateKernelConfig(data_type="float32", default_number_float='float32') + ast = ps.create_kernel(up, config=config) + code = ps.get_code_str(ast) + + assert "1.5f" in code