diff --git a/pystencils/gpu/indexing.py b/pystencils/gpu/indexing.py index 87893af144b53afa5c6e68f2b4ffee7e755361d6..a52cf2a3ff536836c243f1da5dae81f708210f6a 100644 --- a/pystencils/gpu/indexing.py +++ b/pystencils/gpu/indexing.py @@ -398,7 +398,8 @@ def _loop_ctr_assignments(loop_counter_symbols, coordinates, iteration_space): loop_ctr_assignments = [] for loop_counter, coordinate, iter_slice in zip(loop_counter_symbols, coordinates, iteration_space): if isinstance(iter_slice, slice) and iter_slice.step > 1: - loop_ctr_assignments.append(SympyAssignment(loop_counter, coordinate * iter_slice.step - iter_slice.start)) + offset = (iter_slice.step * iter_slice.start) - iter_slice.start + loop_ctr_assignments.append(SympyAssignment(loop_counter, coordinate * iter_slice.step - offset)) elif iter_slice.start == iter_slice.stop: loop_ctr_assignments.append(SympyAssignment(loop_counter, 0)) else: diff --git a/pystencils_tests/test_modulo.py b/pystencils_tests/test_modulo.py index 5a32acf5c3763e1f427e136ab12034b8c817c426..66a3772c477808b08101cae930f4847fc55b4732 100644 --- a/pystencils_tests/test_modulo.py +++ b/pystencils_tests/test_modulo.py @@ -1,16 +1,23 @@ import pytest +import numpy as np import sympy as sp import pystencils as ps from pystencils.astnodes import LoopOverCoordinate, Conditional, Block, SympyAssignment +SLICE_LIST = [False, + ps.make_slice[1:-1:2, 1:-1:2], + ps.make_slice[2:-1:2, 4:-1:7], + ps.make_slice[4:-1:2, 5:-1:2], + ps.make_slice[3:-1:4, 7:-1:3]] + @pytest.mark.parametrize('target', [ps.Target.CPU, ps.Target.GPU]) -@pytest.mark.parametrize('iteration_slice', [False, True]) +@pytest.mark.parametrize('iteration_slice', SLICE_LIST) def test_mod(target, iteration_slice): if target == ps.Target.GPU: pytest.importorskip("cupy") - dh = ps.create_data_handling(domain_size=(5, 5), periodicity=True, default_target=target) + dh = ps.create_data_handling(domain_size=(51, 51), periodicity=True, default_target=target) loop_ctrs = [LoopOverCoordinate.get_loop_counter_symbol(i) for i in range(dh.dim)] cond = [sp.Eq(sp.Mod(loop_ctrs[i], 2), 1) for i in range(dh.dim)] @@ -20,7 +27,6 @@ def test_mod(target, iteration_slice): eq_list = [SympyAssignment(field.center, 1.0)] if iteration_slice: - iteration_slice = ps.make_slice[1:-1:2, 1:-1:2] config = ps.CreateKernelConfig(target=dh.default_target, iteration_slice=iteration_slice) assign = eq_list else: @@ -41,9 +47,5 @@ def test_mod(target, iteration_slice): result = dh.gather_array(field.name, ghost_layers=True) - for x in range(result.shape[0]): - for y in range(result.shape[1]): - if x % 2 == 1 and y % 2 == 1: - assert result[x, y] == 1.0 - else: - assert result[x, y] == 0.0 + assert np.all(result[iteration_slice] == 1.0) +