misc: add test and misc tweaks

Author: mloubout

devito/passes/clusters/misc.py

@@ -99,12 +99,10 @@ def callback(self, clusters, prefix):
 
             # If `c` is made of scalar expressions within guards, then we must keep
             # it close to the adjacent Clusters for correctness
-            if c.is_scalar and c.guards:
-                items = processed
+            if c.is_scalar and c.guards and ispace:
+                processed.append(c.rebuild(ispace=ispace, properties=properties))
             else:
-                items = lifted
-
-            items.append(c.rebuild(ispace=ispace, properties=properties))
+                lifted.append(c.rebuild(ispace=ispace, properties=properties))
 
         return lifted + processed
 

devito/symbolics/search.py

@@ -96,15 +96,18 @@ def visit_preorder_first_hit(self, expr: Expression) -> Iterator[Expression]:
 
 
 def search(exprs: Expression | Iterable[Expression],
-           query: type | Callable[[Any], bool],
+           query: type | tuple[type, ...] | Callable[[Any], bool],
            mode: Mode = 'unique',
            visit: Literal['dfs', 'bfs', 'bfs_first_hit'] = 'dfs',
            deep: bool = False) -> List | set[Expression]:
     """Interface to Search."""
 
     assert mode in ('all', 'unique'), "Unknown mode"
 
-    Q = (lambda obj: isinstance(obj, query)) if isinstance(query, type) else query
+    if isinstance(query, (type, tuple)):
+        Q = lambda obj: isinstance(obj, query)
+    else:
+        Q = query
 
     # Search doesn't actually use a BFS (rather, a preorder DFS), but the terminology
     # is retained in this function's parameters for backwards compatibility

tests/test_dse.py

@@ -2779,7 +2779,7 @@ def test_alias_with_conditional(self):
         cond = FindNodes(Conditional).visit(op)
         assert len(cond) == 3
 
-        # # Each guard should have its own alias for cos(time/ctf)
+        # Each guard should have its own alias for cos(time/ctf)
         scalars = [i for i in FindSymbols().visit(op) if isinstance(i, Temp)]
         assert len(scalars) == 2
 
@@ -2800,8 +2800,76 @@ def test_scalar_alias_interp(self):
 
         op.apply(time_M=3)
 
-        assert np.isclose(norm(f), 254292.75, atol=1e-1, rtol=0)
-        assert np.isclose(norm(s), 191.44644, atol=1e-1, rtol=0)
+        assert np.isclose(norm(f), 254292.75, atol=0, rtol=1e-5)
+        assert np.isclose(norm(s), 191.44644, atol=0, rtol=1e-4)
+
+    def test_scalar_with_cond_access(self):
+        grid = Grid((11, 11))
+        time = grid.time_dim
+
+        u = TimeFunction(name='u', grid=grid, time_order=2, space_order=2)
+
+        ct = ConditionalDimension(name='ct3', parent=time, condition=Ge(time, 2))
+        ct2 = ConditionalDimension(name='ct2', parent=time, factor=4)
+
+        f1 = TimeFunction(name='f1', grid=grid, save=10, time_order=0,
+                          dimensions=(ct,), time_dim=ct, shape=(10,))
+        f1.data[:] = np.arange(10)
+
+        eq0 = Eq(u.forward, u + cos(f1))
+        eq1 = Eq(u.forward, u.forward + sin(time), implicit_dims=ct2)
+        eq2 = Eq(u.forward, u.forward - sin(f1))
+
+        op = Operator([eq0, eq1, eq2])
+
+        cond = FindNodes(Conditional).visit(op)
+        assert len(cond) == 3
+
+        # # Each guard should have its own alias for cos/sin(f1[time-2])
+        scalars = [i for i in FindSymbols().visit(op) if isinstance(i, Temp)]
+        assert len(scalars) == 3
+
+        assert_structure(
+            op,
+            ['t', 't,x,y', 't,x,y', 't,x,y'],
+            'txyxyxy'
+        )
+
+        # This would segfault without the right placement of the alias
+        op.apply(time_M=12)
+
+    def test_scalar_with_cond_tinvariant(self):
+        grid = Grid((10, 10))
+        time = grid.time_dim
+        dt = time.spacing
+
+        u = TimeFunction(name='u', grid=grid, time_order=2, space_order=2)
+
+        ct = ConditionalDimension(name='ct', parent=time, factor=2)
+
+        eq0 = Eq(u.forward, u / dt + 1)
+        eq1 = Eq(u.forward, u.forward + 1/dt**2, implicit_dims=ct)
+
+        op = Operator([eq0, eq1])
+        op(time=5, dt=1)
+
+        cond = FindNodes(Conditional).visit(op)
+        assert len(cond) == 1
+        # One for each 1/dt 1/dt**2
+        scalars = [i for i in FindSymbols().visit(op) if isinstance(i, Temp)]
+        assert len(scalars) == 2
+
+        assert_structure(
+            op,
+            ['t,x,y', 't', 't,x,y'],
+            'txyxy'
+        )
+
+        # Both aliases should be hoisted outside the time loop
+        assert str(body0(op).body[0]) == 'const float r0 = 1.0F/dt;'
+        assert not body0(op).body[0].ispace
+        assert str(body0(op).body[1]) == 'const float r1 = 1.0F/(dt*dt);'
+        assert not body0(op).body[1].ispace
 
 
 class TestIsoAcoustic: