Recurse on stride when inlining (#733)

* Recurse on stride when inlining * Add issue link to test_inline_stride Co-authored-by: Andreas Klöckner <[email protected]>
inducer · Jan 12, 2023 · 4e49be7 · 4e49be7
1 parent c072a86
commit 4e49be7
Show file tree

Hide file tree

Showing 2 changed files with 27 additions and 1 deletion.
diff --git a/loopy/transform/callable.py b/loopy/transform/callable.py
@@ -149,7 +149,7 @@ def map_subscript(self, expr, expn_state):
                 flatten_index += idx*caller_arg.dim_tags[i].stride
 
             flatten_index += sum(
-                idx * tag.stride
+                idx * self.rec(tag.stride, expn_state)
                 for idx, tag in zip(self.rec(expr.index_tuple, expn_state),
                                     callee_arg.dim_tags))
 

diff --git a/test/test_callables.py b/test/test_callables.py
@@ -1404,6 +1404,32 @@ def test_inline_deps(ctx_factory):
     assert np.array_equal(a_dev.get(), np.arange(4))
 
 
+def test_inline_stride():
+    # https://github.com/inducer/loopy/issues/728
+    child_knl = lp.make_function(
+            [],
+            """
+            g[0] = 2*e[0] + 3*f[0] {id=a}
+            g[1] = 2*e[1] + 3*f[1] {dep=a}
+            """, name="linear_combo")
+    parent_knl = lp.make_kernel(
+            ["{[j]:0<=j<n}", "{[i]:0<=i<n}"],
+            """
+            [i]: z[i, j] = linear_combo([i]: x[i, j], [i]: y[i,j])
+            """,
+            kernel_data=[
+                lp.GlobalArg(
+                    name="x, y, z",
+                    dtype=np.float64,
+                    shape=("n", "n")),
+                ...],
+            assumptions="n>=1",
+            )
+    knl = lp.merge([parent_knl, child_knl])
+    knl = lp.inline_callable_kernel(knl, "linear_combo")
+    lp.generate_code_v2(knl).device_code()
+
+
 if __name__ == "__main__":
     if len(sys.argv) > 1:
         exec(sys.argv[1])