From 66fc865c0ad84aaf6364a9d465545bd9388804fd Mon Sep 17 00:00:00 2001
From: Andreas Kloeckner <inform@tiker.net>
Date: Sun, 23 Sep 2012 14:54:25 -0500
Subject: [PATCH] Fixes to ILP write race avoidance.

---
 loopy/preprocess.py |  8 ++++++--
 test/test_loopy.py  | 18 +++++++-----------
 2 files changed, 13 insertions(+), 13 deletions(-)

diff --git a/loopy/preprocess.py b/loopy/preprocess.py
index a1420d7bc..2bdf9a115 100644
--- a/loopy/preprocess.py
+++ b/loopy/preprocess.py
@@ -313,13 +313,17 @@ class ExtraInameIndexInserter(IdentityMapper):
         self.var_to_new_inames = var_to_new_inames
 
     def map_subscript(self, expr):
-        res = IdentityMapper.map_subscript(self, expr)
         try:
             new_idx = self.var_to_new_inames[expr.aggregate.name]
         except KeyError:
             return IdentityMapper.map_subscript(self, expr)
         else:
-            return res.aggregate[res.index + new_idx]
+            index = expr.index
+            if not isinstance(index, tuple):
+                index = (index,)
+            index = tuple(self.rec(i) for i in index)
+
+            return expr.aggregate[index + new_idx]
 
     def map_variable(self, expr):
         try:
diff --git a/test/test_loopy.py b/test/test_loopy.py
index e748e7cc9..0391cf96e 100644
--- a/test/test_loopy.py
+++ b/test/test_loopy.py
@@ -707,25 +707,23 @@ def test_ilp_write_race_detection_global(ctx_factory):
 
 
 
-def test_ilp_write_race_detection_local(ctx_factory):
+def test_ilp_write_race_avoidance_local(ctx_factory):
     ctx = ctx_factory()
 
     knl = lp.make_kernel(ctx.devices[0],
-            "{[i,j]: 0<=i,j<16 }",
+            "{[i,j]: 0<=i<16 and 0<=j<17 }",
             [
                 "[i:l.0, j:ilp] <> a[i] = 5+i+j",
                 ],
             [])
 
-    from loopy.check import WriteRaceConditionError
-    import pytest
-    with pytest.raises(WriteRaceConditionError):
-        list(lp.generate_loop_schedules(knl))
+    for k in lp.generate_loop_schedules(knl):
+        assert k.temporary_variables["a"].shape == (16,17)
 
 
 
 
-def test_ilp_write_race_detection_private(ctx_factory):
+def test_ilp_write_race_avoidance_private(ctx_factory):
     ctx = ctx_factory()
 
     knl = lp.make_kernel(ctx.devices[0],
@@ -735,10 +733,8 @@ def test_ilp_write_race_detection_private(ctx_factory):
                 ],
             [])
 
-    from loopy.check import WriteRaceConditionError
-    import pytest
-    with pytest.raises(WriteRaceConditionError):
-        list(lp.generate_loop_schedules(knl))
+    for k in lp.generate_loop_schedules(knl):
+        assert k.temporary_variables["a"].shape == (16,)
 
 # }}}
 
-- 
GitLab