diff --git a/doc/tutorial.rst b/doc/tutorial.rst index 3c85060dacf03b52f6e0b1faf05ad4697b6a5d07..8ca56619b562e5cffa71974c4ee9f386edfc696b 100644 --- a/doc/tutorial.rst +++ b/doc/tutorial.rst @@ -1220,7 +1220,7 @@ should call :func:`loopy.get_one_scheduled_kernel`: 2: RETURN FROM KERNEL rotate_v2 3: ... gbarrier 4: CALL KERNEL rotate_v2_0(extra_args=[], extra_inames=[]) - 5: arr[((1 + i_inner + i_outer*16) % n)] = tmp {id=rotate} + 5: arr[(1 + i_inner + i_outer*16) % n] = tmp {id=rotate} 6: RETURN FROM KERNEL rotate_v2_0 --------------------------------------------------------------------------- @@ -1260,7 +1260,7 @@ put those instructions into the schedule. 4: ... gbarrier 5: CALL KERNEL rotate_v2_0(extra_args=['tmp_save_slot'], extra_inames=[]) 6: tmp = tmp_save_slot[tmp_reload_hw_dim_0_rotate_v2_0, tmp_reload_hw_dim_1_rotate_v2_0] {id=tmp.reload} - 7: arr[((1 + i_inner + i_outer*16) % n)] = tmp {id=rotate} + 7: arr[(1 + i_inner + i_outer*16) % n] = tmp {id=rotate} 8: RETURN FROM KERNEL rotate_v2_0 --------------------------------------------------------------------------- diff --git a/loopy/check.py b/loopy/check.py index 5fb7a2fb928a26fd44fa0ff22598e71646853de6..e11aafab093c2f74875caa372dacbc9306fd9648 100644 --- a/loopy/check.py +++ b/loopy/check.py @@ -29,6 +29,9 @@ from islpy import dim_type import islpy as isl from loopy.symbolic import WalkMapper from loopy.diagnostic import LoopyError, WriteRaceConditionWarning, warn_with_kernel +from loopy.type_inference import TypeInferenceMapper +from loopy.kernel.instruction import (MultiAssignmentBase, CallInstruction, + CInstruction, _DataObliviousInstruction) import logging logger = logging.getLogger(__name__) @@ -66,6 +69,32 @@ def check_identifiers_in_subst_rules(knl): VALID_NOSYNC_SCOPES = frozenset(["local", "global", "any"]) +class SubscriptIndicesIsIntChecker(TypeInferenceMapper): + def map_subscript(self, expr): + for idx in expr.index_tuple: + if not self.rec(idx)[0].is_integral(): + raise LoopyError("Non integral array indices obtained in" + " {}.".format(expr)) + + return self.rec(expr.aggregate) + + +def check_for_integer_subscript_indices(kernel): + from pymbolic.primitives import Subscript + idx_int_checker = SubscriptIndicesIsIntChecker(kernel) + for insn in kernel.instructions: + if isinstance(insn, MultiAssignmentBase): + idx_int_checker(insn.expression, return_tuple=isinstance(insn, + CallInstruction), return_dtype_set=True) + [idx_int_checker(assignee) for assignee in insn.assignees if + isinstance(assignee, Subscript)] + elif isinstance(insn, (CInstruction, _DataObliviousInstruction)): + pass + else: + raise NotImplementedError("Unknown insn type %s." % ( + type(insn).__name__)) + + def check_insn_attributes(kernel): all_insn_ids = set(insn.id for insn in kernel.instructions) @@ -620,6 +649,7 @@ def pre_schedule_checks(kernel): try: logger.debug("%s: pre-schedule check: start" % kernel.name) + check_for_integer_subscript_indices(kernel) check_for_duplicate_insn_ids(kernel) check_for_orphaned_user_hardware_axes(kernel) check_for_double_use_of_hw_axes(kernel) diff --git a/test/test_loopy.py b/test/test_loopy.py index 119d57adf2c850eba3bb6ad5df3c0a8d0644b70c..d506258c250ddf790a7e74b9e01fe0d4967d0850 100644 --- a/test/test_loopy.py +++ b/test/test_loopy.py @@ -3015,6 +3015,19 @@ def test_array_arg_extra_kwargs_persis_hash(): assert key_builder(a) != key_builder(not_a) +def test_non_integral_array_idx_raises(): + knl = lp.make_kernel( + "{[i, j]: 0<=i<=4 and 0<=j<16}", + """ + out[j] = 0 {id=init} + out[i] = a[1.94**i-1] {dep=init} + """, [lp.GlobalArg('a', np.float64), '...']) + + from loopy.diagnostic import LoopyError + with pytest.raises(LoopyError): + print(lp.generate_code_v2(knl).device_code()) + + if __name__ == "__main__": if len(sys.argv) > 1: exec(sys.argv[1])