compiler: improve lifting processing to avoid aliases missplacement

Author: mloubout

conftest.py

@@ -479,3 +479,9 @@ def check_array(array, exp_halo, exp_shape, rotate=False):
 
     assert tuple(array.halo) == exp_halo
     assert tuple(shape) == tuple(exp_shape)
+
+
+# Main body in Operator IET, depending on ISA
+def body0(op):
+    bidx = 0 if 'sse' not in configuration['platform'].known_isas else 1
+    return op.body.body[bidx]

devito/passes/clusters/aliases.py

@@ -127,6 +127,8 @@ def _aliases_from_clusters(self, cgroup, exclude, meta):
         for mapper in self._generate(cgroup, exclude):
             # Clusters -> AliasList
             found = collect(mapper.extracted, meta.ispace, self.opt_minstorage)
+            if not found:
+                continue
             exprs, aliases = self._choose(found, cgroup, mapper)
 
             # AliasList -> Schedule
@@ -147,7 +149,7 @@ def _aliases_from_clusters(self, cgroup, exclude, meta):
         # Schedule -> [Clusters]_k
         processed, subs = lower_schedule(schedule, meta, self.sregistry,
                                          self.opt_ftemps, self.opt_min_dtype,
-                                         self.opt_minmem)
+                                         self.opt_minmem, nclusters=len(cgroup))
 
         # [Clusters]_k -> [Clusters]_k (optimization)
         if self.opt_multisubdomain:
@@ -272,7 +274,6 @@ def _do_generate(self, exprs, exclude, cbk_search, cbk_compose=None):
                     free_symbols = i.free_symbols
                 if {a.function for a in free_symbols} & exclude:
                     continue
-
                 mapper.add(i, make, terms)
 
         return mapper
@@ -853,7 +854,7 @@ def optimize_schedule_rotations(schedule, sregistry):
 
 
 def lower_schedule(schedule, meta, sregistry, opt_ftemps, opt_min_dtype,
-                   opt_minmem):
+                   opt_minmem, nclusters=1):
     """
     Turn a Schedule into a sequence of Clusters.
     """
@@ -929,20 +930,21 @@ def lower_schedule(schedule, meta, sregistry, opt_ftemps, opt_min_dtype,
             # Degenerate case: scalar expression
             assert writeto.size == 0
 
-            guards = None
             is_cond = any(isinstance(d, (SubsamplingFactor, ConditionalDimension))
                           for d in pivot.free_symbols)
-            if meta.guards and is_cond:
+            if meta.guards and is_cond and nclusters > 1:
                 # Scalar alias that depends on a guard, unsafe to lift out of the guard
                 # Do not alias
                 expression = None
                 callback = lambda idx: uxreplace(pivot, subs)  # noqa: B023
             else:
                 dtype = sympy_dtype(pivot, base=meta.dtype, smin=opt_min_dtype)
-                obj = Temp(name=name, dtype=dtype)
+                obj = Temp(name=name, dtype=dtype, is_const=True)
                 expression = Eq(obj, uxreplace(pivot, subs))
 
                 callback = lambda idx: obj  # noqa: B023
+                # Only keep the guard if there is no cross-cluster reuse of the scalar
+                guards = meta.guards if nclusters == 1 else None
 
         # Create the substitution rules for the aliasing expressions
         subs.update({

devito/passes/clusters/misc.py

@@ -102,7 +102,14 @@ def callback(self, clusters, prefix):
             # unless the guard is for an outer dimension
             guards = {} if c.is_scalar and not (prefix[:-1] and c.guards) else c.guards
 
-            lifted.append(c.rebuild(ispace=ispace, properties=properties, guards=guards))
+            _lifted = c.rebuild(ispace=ispace, properties=properties, guards=guards)
+            if guards and clusters[max(n-1, 0)].guards != guards and _lifted.is_scalar:
+                # Heuristic: if the lifted Cluster has different guards than the
+                # previous one, then we are likely to end up with a separate
+                # Cluster, hence give up on lifting
+                processed.append(_lifted)
+            else:
+                lifted.append(_lifted)
 
         return lifted + processed
 

examples/performance/01_gpu.ipynb

@@ -142,8 +142,13 @@
      "name": "stderr",
      "output_type": "stream",
      "text": [
-      "NUMA domain count autodetection failed, assuming 1\n",
-      "Operator `Kernel` ran in 0.01 s\n",
+      "NUMA domain count autodetection failed, assuming 1\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
       "Operator `Kernel` ran in 0.01 s\n"
      ]
     }
@@ -292,9 +297,9 @@
       "  const int x_stride0 = x_fsz0*y_fsz0;\n",
       "  const int y_stride0 = y_fsz0;\n",
       "\n",
-      "  float r0 = 1.0F/dt;\n",
-      "  float r1 = 1.0F/(h_x*h_x);\n",
-      "  float r2 = 1.0F/(h_y*h_y);\n",
+      "  const float r0 = 1.0F/dt;\n",
+      "  const float r1 = 1.0F/(h_x*h_x);\n",
+      "  const float r2 = 1.0F/(h_y*h_y);\n",
       "\n",
       "  for (int time = time_m; time <= time_M; time += 1)\n",
       "  {\n",
@@ -340,7 +345,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.13.5"
+   "version": "3.13.11"
   }
  },
  "nbformat": 4,

tests/test_dse.py

@@ -5,8 +5,8 @@
 from sympy import Mul  # noqa
 
 from conftest import (  # noqa
-    _R, EVAL, assert_blocking, assert_structure, check_array, get_arrays, get_params,
-    skipif
+    _R, EVAL, assert_blocking, assert_structure, body0, check_array, get_arrays,
+    get_params, skipif
 )
 from devito import (  # noqa
     NODE, Abs, ConditionalDimension, Constant, DefaultDimension, Derivative, Dimension,
@@ -348,8 +348,8 @@ def test_scalar_cond(self):
         trees = retrieve_iteration_tree(op)
 
         assert len(trees) == 3
-        assert_structure(op, ['t', 't,x,y', 't,x,y'], 'txyxy')
-        assert trees[0].dimensions == [time]
+        assert_structure(op, ['t,x,y', 't', 't,x,y'], 'txyxy')
+        assert trees[1].dimensions == [time]
 
 
 class TestAliases:
@@ -2552,6 +2552,7 @@ def test_invariants_with_conditional(self):
         eqn = Eq(u, u - (cos(time_sub * factor * f) * uf))
 
         op = Operator(eqn, opt='advanced')
+
         assert_structure(op, ['t', 't,fd', 't,fd,x,y'], 't,fd,x,y')
         # Make sure it compiles
         _ = op.cfunction
@@ -2700,10 +2701,12 @@ def test_split_cond(self):
         eq2 = Eq(u.forward, u.forward + cos(time), implicit_dims=ct)
 
         op = Operator([eq0, eq1, eq2])
+        op(time=5)
+
         cond = FindNodes(Conditional).visit(op)
         assert len(cond) == 3
         # The alias should have been lifted out of the condition
-        assert 'float r0 = cos(time);' in str(op.body.body[0])
+        assert 'float r0 = cos(time);' in str(body0(op))
         scalars = [i for i in FindSymbols().visit(op) if isinstance(i, Temp)]
         assert len(scalars) == 1
 
@@ -2721,10 +2724,12 @@ def test_split_cond_multi_alias(self):
         eq2 = Eq(u.forward, u.forward + cos(time) - sin(time), implicit_dims=ct)
 
         op = Operator([eq0, eq1, eq2])
+        op(time=5)
+
         cond = FindNodes(Conditional).visit(op)
         assert len(cond) == 3
         # The alias should have been lifted out of the condition
-        assert 'float r3 = cos(time);' in str(op.body.body[0])
+        assert 'float r3 = cos(time);' in str(body0(op))
         scalars = [i for i in FindSymbols().visit(op) if isinstance(i, Temp)]
         assert len(scalars) == 5
 
@@ -2743,6 +2748,7 @@ def test_multi_cond_no_split(self):
         eq2 = Eq(u.forward, u.forward - sin(time), implicit_dims=ct)
 
         op = Operator([eq0, eq1, eq2])
+        op(time=5)
 
         assert_structure(
             op,
@@ -2751,7 +2757,7 @@ def test_multi_cond_no_split(self):
         )
 
         scalars = [i for i in FindSymbols().visit(op) if isinstance(i, Temp)]
-        assert len(scalars) == 4
+        assert len(scalars) == 3
 
     def test_alias_with_conditional(self):
         grid = Grid((11, 11))
@@ -2767,6 +2773,8 @@ def test_alias_with_conditional(self):
         eq2 = Eq(u.forward, u.forward + cos(ct), implicit_dims=ct)
 
         op = Operator([eq0, eq1, eq2])
+        op(time=5)
+
         cond = FindNodes(Conditional).visit(op)
         assert len(cond) == 3
 

tests/test_mpi.py

@@ -4,7 +4,7 @@
 import pytest
 from test_dse import TestTTI
 
-from conftest import _R, assert_blocking, assert_structure
+from conftest import _R, assert_blocking, assert_structure, body0
 from devito import (
     NODE, Buffer, ConditionalDimension, Constant, CustomDimension, DefaultDimension,
     Dimension, Eq, Function, Grid, Inc, Ne, Operator, PrecomputedSparseFunction,
@@ -24,12 +24,6 @@
 from examples.seismic.acoustic import acoustic_setup
 
 
-# Main body in Operator IET, depending on ISA
-def body0(op):
-    bidx = 0 if 'sse' not in configuration['platform'].known_isas else 1
-    return op.body.body[bidx]
-
-
 class TestDistributor:
 
     @pytest.mark.parallel(mode=[2, 4])

tests/test_symbolics.py

@@ -58,8 +58,8 @@ def test_func_of_indices():
 
 
 @pytest.mark.parametrize('dtype,expected', [
-    (np.float32, "float r0 = 1.0F/h_x;"),
-    (np.float64, "double r0 = 1.0/h_x;")
+    (np.float32, "const float r0 = 1.0F/h_x;"),
+    (np.float64, "const double r0 = 1.0/h_x;")
 ])
 def test_floatification_issue_1627(dtype, expected):
     """