diff --git a/test/test_transform.py b/test/test_transform.py index ee18c81141071bd6f0999085a170e48ca7551b7b..bb7b34cf220c692934a32415b002782f618a3d18 100644 --- a/test/test_transform.py +++ b/test/test_transform.py @@ -52,6 +52,24 @@ __all__ = [ from loopy.version import LOOPY_USE_LANGUAGE_VERSION_2018_2 # noqa +# {{{ ContainsFloorDiv + +class ContainsFloorDiv(lp.symbolic.CombineMapper): + def combine(self, values): + return any(values) + + def map_floor_div(self, expr): + return True + + def map_variable(self, expr): + return False + + def map_constant(self, expr): + return False + +# }}} + + @pytest.mark.parametrize("fix_parameters", (True, False)) def test_chunk_iname(ctx_factory, fix_parameters): ctx = ctx_factory() @@ -1260,28 +1278,44 @@ def test_privatize_with_nonzero_lbound(ctx_factory): np.testing.assert_allclose(out.get()[10:14], np.arange(10, 14)) -def test_simplify_indices(ctx_factory): +def test_simplify_indices_when_inlining(ctx_factory): ctx = ctx_factory() + twice = lp.make_function( + "{[i, j]: 0<=i<10 and 0<=j<4}", + """ + y[i,j] = 2*x[i,j] + """, name="zerozerozeroonezeroify") + knl = lp.make_kernel( - "{[i]: 0<=i<10}", + "{:}", """ - Y[i] = X[10*(i//10) + i] + Y[:,:] = zerozerozeroonezeroify(X[:,:]) """, [lp.GlobalArg("X,Y", - shape=(10,), + shape=(10, 4), dtype=np.float64)]) - class ContainsFloorDiv(lp.symbolic.CombineMapper): - def combine(self, values): - return any(values) + knl = lp.merge([knl, twice]) + inlined_knl = lp.inline_callable_kernel(knl, "zerozerozeroonezeroify") + contains_floordiv = ContainsFloorDiv() + + print(inlined_knl) - def map_floor_div(self, expr): - return True + assert all(not contains_floordiv(insn.expression) + for insn in inlined_knl.default_entrypoint.instructions + if isinstance(insn, lp.MultiAssignmentBase)) - def map_variable(self, expr): - return False + lp.auto_test_vs_ref(knl, ctx, inlined_knl) - def map_constant(self, expr): - return False + +def test_simplify_indices(ctx_factory): + ctx = ctx_factory() + knl = lp.make_kernel( + "{[i]: 0<=i<10}", + """ + Y[i] = X[10*(i//10) + i] + """, [lp.GlobalArg("X,Y", + shape=(10,), + dtype=np.float64)]) simplified_knl = lp.simplify_indices(knl) contains_floordiv = ContainsFloorDiv()