From 6c82b3785024eee0c8d41e5752ae28a90f02c28d Mon Sep 17 00:00:00 2001
From: Andreas Kloeckner <inform@tiker.net>
Date: Fri, 27 Oct 2017 19:31:55 +0200
Subject: [PATCH 1/2] Make sure add_prefetch also works in subexpressions of
 assignment LHSs

---
 loopy/transform/subst.py |  9 ++++-----
 test/test_loopy.py       | 25 +++++++++++++++++++++++++
 2 files changed, 29 insertions(+), 5 deletions(-)

diff --git a/loopy/transform/subst.py b/loopy/transform/subst.py
index 79ceff9fd..a681afe06 100644
--- a/loopy/transform/subst.py
+++ b/loopy/transform/subst.py
@@ -1,6 +1,4 @@
-from __future__ import division
-from __future__ import absolute_import
-import six
+from __future__ import division, absolute_import
 
 __copyright__ = "Copyright (C) 2012 Andreas Kloeckner"
 
@@ -24,6 +22,7 @@ OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
 THE SOFTWARE.
 """
 
+import six
 
 from loopy.symbolic import (
         get_dependencies, SubstitutionMapper,
@@ -141,6 +140,7 @@ def extract_subst(kernel, subst_name, template, parameters=()):
     dfmapper = CallbackMapper(gather_exprs, WalkMapper())
 
     for insn in kernel.instructions:
+        dfmapper(insn.assignees)
         dfmapper(insn.expression)
 
     for sr in six.itervalues(kernel.substitutions):
@@ -178,8 +178,7 @@ def extract_subst(kernel, subst_name, template, parameters=()):
     new_insns = []
 
     for insn in kernel.instructions:
-        new_expr = cbmapper(insn.expression)
-        new_insns.append(insn.copy(expression=new_expr))
+        new_insns.append(insn.with_transformed_expressions(cbmapper))
 
     from loopy.kernel.data import SubstitutionRule
     new_substs = {
diff --git a/test/test_loopy.py b/test/test_loopy.py
index f2658ca75..deebc9db3 100644
--- a/test/test_loopy.py
+++ b/test/test_loopy.py
@@ -2678,6 +2678,31 @@ def test_preamble_with_separate_temporaries(ctx_factory):
         queue, data=data.flatten('C'))[1][0], data[offsets[:-1] + 1])
 
 
+def test_add_prefetch_works_in_lhs_index():
+    knl = lp.make_kernel(
+            "{ [n,k,l,k1,l1,k2,l2]: "
+            "start<=n<end and 0<=k,k1,k2<3 and 0<=l,l1,l2<2 }",
+            """
+            for n
+                <> a1_tmp[k,l] = a1[a1_map[n, k],l]
+                a1_tmp[k1,l1] = a1_tmp[k1,l1] + 1
+                a1_out[a1_map[n,k2], l2] = a1_tmp[k2,l2]
+            end
+            """,
+            [
+                lp.GlobalArg("a1,a1_out", None, "ndofs,2"),
+                lp.GlobalArg("a1_map", None, "nelements,3"),
+                ...
+            ])
+
+    knl = lp.add_prefetch(knl, "a1_map", "k")
+
+    print(knl)
+    from loopy.symbolic import get_dependencies
+    for insn in knl.instructions:
+        assert "a1_map" not in get_dependencies(insn.assignees)
+
+
 if __name__ == "__main__":
     if len(sys.argv) > 1:
         exec(sys.argv[1])
-- 
GitLab


From 6cf3e8aabb12263f41d37cff369e041f99cb1517 Mon Sep 17 00:00:00 2001
From: Andreas Kloeckner <inform@tiker.net>
Date: Sat, 28 Oct 2017 16:50:22 +0200
Subject: [PATCH 2/2] Py2 ellipsis fix

---
 test/test_loopy.py | 3 +--
 1 file changed, 1 insertion(+), 2 deletions(-)

diff --git a/test/test_loopy.py b/test/test_loopy.py
index deebc9db3..d0398f216 100644
--- a/test/test_loopy.py
+++ b/test/test_loopy.py
@@ -2692,12 +2692,11 @@ def test_add_prefetch_works_in_lhs_index():
             [
                 lp.GlobalArg("a1,a1_out", None, "ndofs,2"),
                 lp.GlobalArg("a1_map", None, "nelements,3"),
-                ...
+                "..."
             ])
 
     knl = lp.add_prefetch(knl, "a1_map", "k")
 
-    print(knl)
     from loopy.symbolic import get_dependencies
     for insn in knl.instructions:
         assert "a1_map" not in get_dependencies(insn.assignees)
-- 
GitLab