Merge pull request #2637 from devitocodes/cse-dequeue-fix

compiler: fix deque pop order

Author: mloubout

devito/passes/clusters/cse.py

@@ -279,7 +279,8 @@ def choose_element(queue, scheduled):
             first = sorted(tmps, key=lambda i: i.lhs.name).pop(0)
             queue.remove(first)
         else:
-            first = queue.pop()
+            first = sorted(queue, key=lambda i: exprs.index(i)).pop(0)
+            queue.remove(first)
         return first
 
     processed = dag.topological_sort(choose_element)

requirements.txt

@@ -1,5 +1,5 @@
 pip>=9.0.1
-numpy>=2,<2.3.0
+numpy>=2,<2.3.1
 sympy>=1.12.1,<1.15
 psutil>=5.1.0,<8.0
 py-cpuinfo<10

tests/test_cse.py

@@ -246,3 +246,26 @@ def test_advanced_algo(exprs, expected):
 
     assert len(processed) == len(expected)
     assert all(str(i.rhs) == j for i, j in zip(processed, expected))
+
+
+def test_advanced_algo_order():
+    """
+    Test that smartsort/advanced doesn't break equation order.
+    """
+    grid = Grid((3, 3, 3))
+    u = TimeFunction(name="u", grid=grid, space_order=2)
+    v = TimeFunction(name="v", grid=grid, space_order=2)
+
+    eq0 = DummyEq(indexify(diffify(Eq(u.forward, u.dx).evaluate)))
+    eq1 = DummyEq(indexify(diffify(Eq(v, u.dx).evaluate)))
+    eq_b = DummyEq(indexify(diffify(Eq(v.forward, v + u.forward).evaluate)))
+
+    counter = generator()
+    make = lambda _: CTemp(name='r%d' % counter(), dtype=np.float32).indexify()
+    processed = _cse([eq0, eq1, eq_b], make, mode='advanced')
+
+    # Three input equation and 2 CTemps
+    assert len(processed) == 5
+    assert processed[0].lhs.name == 'r1'
+    # eq_b has to be last
+    assert processed[-1] == eq_b