diff --git a/src/pystencils/gpu/indexing.py b/src/pystencils/gpu/indexing.py index 843e77bb87f60ea4c509d4ec8827e41ee204c42f..cecd0642415253c40c6a61fe92a22040d4fc5501 100644 --- a/src/pystencils/gpu/indexing.py +++ b/src/pystencils/gpu/indexing.py @@ -8,7 +8,7 @@ from sympy.core.cache import cacheit from pystencils.astnodes import Block, Conditional, SympyAssignment from pystencils.typing import TypedSymbol, create_type -from pystencils.integer_functions import div_ceil, div_floor +from pystencils.integer_functions import div_ceil, div_floor, int_div from pystencils.sympyextensions import is_integer_sequence, prod @@ -224,6 +224,9 @@ class BlockIndexing(AbstractIndexing): assert len(self._iteration_space) == len(arr_shape), "Iteration space must be equal to the array shape" numeric_iteration_slice = _get_numeric_iteration_slice(self._iteration_space, arr_shape) end = [s.stop if s.stop != 0 else 1 for s in numeric_iteration_slice] + for i, s in enumerate(numeric_iteration_slice): + if s.step and s.step != 1: + end[i] = int_div(end[i], s.step) if self._dim < 4: conditions = [c < e for c, e in zip(self.coordinates, end)]