From be9636d28691ce0b737226d969b7d9796ac38b6f Mon Sep 17 00:00:00 2001
From: Kaushik Kulkarni <15399010+kaushikcfd@users.noreply.github.com>
Date: Sat, 5 Jun 2021 18:27:55 -0500
Subject: [PATCH] [callables] Test kernel call with non 1-step slice (#396)

* test a non1-step slice

* highlight the difference between make_slab and slice's access map

* corrects domain for slices with 'start'

* sharpens error condition

- non-integral step sizes cannot be supported in loopy's domain representation

* tests slices with more initial start expressions

* make_slab: step -> iname_multiplier, disallow -ve multipliers

- make_slab and slices are different, using the same names is confusion. Besides
that, 'step' wasn't a representative name for the input
---
 loopy/isl_helpers.py     | 38 ++++++++++++++----------------
 loopy/kernel/creation.py | 13 +++++++++-
 test/test_callables.py   | 51 ++++++++++++++++++++++++++++++++++++++++
 3 files changed, 80 insertions(+), 22 deletions(-)

diff --git a/loopy/isl_helpers.py b/loopy/isl_helpers.py
index 3e841ba87..d67df1154 100644
--- a/loopy/isl_helpers.py
+++ b/loopy/isl_helpers.py
@@ -60,29 +60,34 @@ def dump_space(ls):
 
 # {{{ make_slab
 
-def make_slab(space, iname, start, stop, step=1):
+def make_slab(space, iname, start, stop, iname_multiplier=1):
     """
     Returns an instance of :class:`islpy._isl.BasicSet`, which satisfies the
-    constraint ``start <= step*iname < stop``.
+    constraint ``start <= iname_multiplier*iname < stop``.
 
     :arg space: An instance of :class:`islpy._isl.Space`.
 
     :arg iname:
+
         Either an instance of :class:`str` as a name of the ``iname`` or a
         tuple of ``(iname_dt, iname_dx)`` indicating the *iname* in the space.
 
     :arg start:
+
         An instance of :class:`int`  or an instance of
         :class:`islpy._isl.Aff` indicating the lower bound of
-        ``step*iname``(inclusive).
+        ``iname_multiplier*iname``(inclusive).
 
     :arg stop:
+
         An instance of :class:`int`  or an instance of
         :class:`islpy._isl.Aff` indicating the upper bound of
-        ``step*iname``.
+        ``iname_multiplier*iname``.
+
+    :arg iname_multiplier:
 
-    :arg step:
-        An instance of :class:`int`.
+        A strictly positive :class:`int` denoting *iname*'s coefficient in the
+        above inequality expression.
     """
     zero = isl.Aff.zero_on_domain(space)
 
@@ -112,25 +117,16 @@ def make_slab(space, iname, start, stop, step=1):
 
     iname_aff = zero.add_coefficient_val(iname_dt, iname_idx, 1)
 
-    if step > 0:
-        result = (isl.BasicSet.universe(space)
-                # start <= step*iname
-                .add_constraint(isl.Constraint.inequality_from_aff(
-                    step*iname_aff - start))
-                # step*iname < stop
-                .add_constraint(isl.Constraint.inequality_from_aff(
-                    stop-1 - step*iname_aff)))
-    elif step < 0:
+    if iname_multiplier > 0:
         result = (isl.BasicSet.universe(space)
-                # start >= (-step)*iname
+                # start <= iname_multiplier*iname
                 .add_constraint(isl.Constraint.inequality_from_aff(
-                    step*iname_aff + start))
-                # (-step)*iname > stop
+                    iname_multiplier*iname_aff - start))
+                # iname_multiplier*iname < stop
                 .add_constraint(isl.Constraint.inequality_from_aff(
-                    -stop-1 - step*iname_aff)))
+                    stop-1 - iname_multiplier*iname_aff)))
     else:
-        # step = 0
-        raise LoopyError("0 step not allowed in make_slab.")
+        raise LoopyError("iname_multiplier must be strictly positive")
 
     return result
 
diff --git a/loopy/kernel/creation.py b/loopy/kernel/creation.py
index 4d1e86ca7..b9cf234c6 100644
--- a/loopy/kernel/creation.py
+++ b/loopy/kernel/creation.py
@@ -1899,6 +1899,8 @@ def normalize_slice_params(slice, dimension_length):
     :arg dimension_length: Length of the axis being sliced.
     """
     from pymbolic.primitives import Slice
+    from numbers import Integral
+
     assert isinstance(slice, Slice)
     start, stop, step = slice.start, slice.stop, slice.step
 
@@ -1924,6 +1926,10 @@ def normalize_slice_params(slice, dimension_length):
 
     # }}}
 
+    if not isinstance(step, Integral):
+        raise LoopyError("Non-integral step sizes lead to non-affine domains =>"
+                         " not supported")
+
     return start, stop, step
 
 
@@ -2063,7 +2069,12 @@ class SliceToInameReplacer(IdentityMapper):
 
             from loopy.isl_helpers import make_slab
             for iname, (start, stop, step) in sar_bounds.items():
-                iname_set = iname_set & make_slab(space, iname, start, stop, step)
+                if step > 0:
+                    iname_set = iname_set & make_slab(space, iname, 0,
+                                                      stop-start, step)
+                else:
+                    iname_set = iname_set & make_slab(space, iname, 0,
+                                                      start-stop, -step)
 
             subarray_ref_domains.append(iname_set)
 
diff --git a/test/test_callables.py b/test/test_callables.py
index 4bc37aeba..c19c7f1d0 100644
--- a/test/test_callables.py
+++ b/test/test_callables.py
@@ -855,6 +855,57 @@ def test_kc_with_floor_div_in_expr(ctx_factory, inline):
     lp.auto_test_vs_ref(knl, ctx, knl)
 
 
+@pytest.mark.parametrize("start", [5, 6, 7])
+@pytest.mark.parametrize("inline", [True, False])
+def test_non1_step_slices(ctx_factory, start, inline):
+    # See https://github.com/inducer/loopy/pull/222#discussion_r645905188
+
+    ctx = ctx_factory()
+    cq = cl.CommandQueue(ctx)
+
+    callee = lp.make_function(
+            "{[i]: 0<=i<n}",
+            """
+            y[i] = i**2
+            """,
+            [lp.ValueArg("n"), ...],
+            name="squared_arange")
+
+    t_unit = lp.make_kernel(
+            "{[i_init, j_init]: 0<=i_init, j_init<40}",
+            f"""
+            X[i_init] = 42
+            X[{start}:40:3] = squared_arange({len(range(start, 40, 3))})
+
+            Y[j_init] = 1729
+            Y[39:{start}:-3] = squared_arange({len(range(39, start, -3))})
+            """,
+            [lp.GlobalArg("X,Y", shape=40)],
+            seq_dependencies=True)
+
+    expected_out1 = 42*np.ones(40, dtype=np.int64)
+    expected_out1[start:40:3] = np.arange(len(range(start, 40, 3)))**2
+
+    expected_out2 = 1729*np.ones(40, dtype=np.int64)
+    expected_out2[39:start:-3] = np.arange(len(range(39, start, -3)))**2
+
+    t_unit = lp.merge([t_unit, callee])
+
+    t_unit = lp.set_options(t_unit, "return_dict")
+
+    # check_bounds isn't smart enough to pass assumptions from caller to callee
+    # during check_bounds
+    t_unit = lp.set_options(t_unit, enforce_array_accesses_within_bounds="no_check")
+
+    if inline:
+        t_unit = lp.inline_callable_kernel(t_unit, "squared_arange")
+
+    evt, out_dict = t_unit(cq)
+
+    np.testing.assert_allclose(out_dict["X"].get(), expected_out1)
+    np.testing.assert_allclose(out_dict["Y"].get(), expected_out2)
+
+
 if __name__ == "__main__":
     if len(sys.argv) > 1:
         exec(sys.argv[1])
-- 
GitLab