diff --git a/docs/source/user_manual/reductions.md b/docs/source/user_manual/reductions.md index 454fd7df4c4f686075c70e45f38bf2d5563ae6af..6af0e8580f215f8b205e3250e65ec482870b4523 100644 --- a/docs/source/user_manual/reductions.md +++ b/docs/source/user_manual/reductions.md @@ -56,21 +56,18 @@ The following reduction assignment classes are available in pystencils: * `MaxReductionAssignment`: Finds maximum element :::{note} -Alternatívely, you can also make use of the `reduction_assignment` or `reduction_assignment_from_str` functions +Alternatívely, you can also make use of the `reduction_assignment` function to specify reduction assignments: ::: ```{code-cell} ipython3 -from pystencils.sympyextensions import reduction_assignment, reduction_assignment_from_str +from pystencils.sympyextensions import reduction_assignment from pystencils.sympyextensions.reduction import ReductionOp assign_sum = reduction_assignment(r, ReductionOp.Add, x.center()) - -assign_sum = reduction_assignment_from_str(r, "+", x.center()) ``` -For other reduction operations, the following enums can be passed to `reduction_assignment` -or the corresponding strings can be passed to `reduction_assignment_from_str`. +For other reduction operations, the following enums can be passed to `reduction_assignment`. ```{code-cell} python3 class ReductionOp(Enum): diff --git a/src/pystencils/sympyextensions/__init__.py b/src/pystencils/sympyextensions/__init__.py index bd0fa1fe9d5459f842ec4aca52bf13de63ab7834..c575feeb3d973ba77d3ecaeb461730dd2b427fef 100644 --- a/src/pystencils/sympyextensions/__init__.py +++ b/src/pystencils/sympyextensions/__init__.py @@ -1,7 +1,7 @@ from .astnodes import ConditionalFieldAccess from .typed_sympy import TypedSymbol, CastFunc, tcast, DynamicType from .pointers import mem_acc -from .reduction import reduction_assignment, reduction_assignment_from_str, ReductionOp +from .reduction import reduction_assignment, ReductionOp from .math import ( prod, @@ -35,7 +35,6 @@ from .math import ( __all__ = [ "ConditionalFieldAccess", "reduction_assignment", - "reduction_assignment_from_str", "ReductionOp", "TypedSymbol", "CastFunc", diff --git a/tests/kernelcreation/test_reduction.py b/tests/kernelcreation/test_reduction.py index cd1710cf5ffe8e5ced135867be2c9e947c7e6b64..1fb8efc818797d37717277c0a9a74e3734e3a556 100644 --- a/tests/kernelcreation/test_reduction.py +++ b/tests/kernelcreation/test_reduction.py @@ -2,17 +2,17 @@ import pytest import numpy as np import pystencils as ps -from pystencils.sympyextensions import reduction_assignment_from_str +from pystencils.sympyextensions import ReductionOp, reduction_assignment INIT_W = 5 INIT_ARR = 2 SIZE = 15 SOLUTION = { - "+": INIT_W + INIT_ARR * SIZE, - "-": INIT_W - INIT_ARR * SIZE, - "*": INIT_W * INIT_ARR**SIZE, - "min": min(INIT_W, INIT_ARR), - "max": max(INIT_W, INIT_ARR), + ReductionOp.Add: INIT_W + INIT_ARR * SIZE, + ReductionOp.Sub: INIT_W - INIT_ARR * SIZE, + ReductionOp.Mul: INIT_W * INIT_ARR**SIZE, + ReductionOp.Min: min(INIT_W, INIT_ARR), + ReductionOp.Max: max(INIT_W, INIT_ARR), } @@ -21,14 +21,14 @@ def get_reduction_assign_ast(dtype, op, config): x = ps.fields(f"x: {dtype}[1d]") w = ps.TypedSymbol("w", dtype) - red_assign = reduction_assignment_from_str(w, op, x.center()) + red_assign = reduction_assignment(w, op, x.center()) return ps.create_kernel([red_assign], config, default_dtype=dtype) @pytest.mark.parametrize("instruction_set", ["sse", "avx"]) @pytest.mark.parametrize("dtype", ["float64", "float32"]) -@pytest.mark.parametrize("op", ["+", "-", "*", "min", "max"]) +@pytest.mark.parametrize("op", [ReductionOp.Add, ReductionOp.Sub, ReductionOp.Mul, ReductionOp.Min, ReductionOp.Max]) def test_reduction_cpu(instruction_set, dtype, op): vectorize_info = { "instruction_set": instruction_set, @@ -51,7 +51,7 @@ def test_reduction_cpu(instruction_set, dtype, op): @pytest.mark.parametrize("dtype", ["float64", "float32"]) -@pytest.mark.parametrize("op", ["+", "-", "*", "min", "max"]) +@pytest.mark.parametrize("op", [ReductionOp.Add, ReductionOp.Sub, ReductionOp.Mul, ReductionOp.Min, ReductionOp.Max]) @pytest.mark.parametrize("assume_warp_aligned_block_size", [True, False]) @pytest.mark.parametrize("use_block_fitting", [True, False]) def test_reduction_gpu(