compiler: improve lifting processing to avoid aliases missplacement

Author: mloubout

devito/passes/clusters/aliases.py

@@ -132,7 +132,8 @@ def _aliases_from_clusters(self, cgroup, exclude, meta):
             # AliasList -> Schedule
             schedule = lower_aliases(aliases, meta, self.opt_maxpar)
 
-            variants.append(Variant(schedule, exprs))
+            if schedule:
+                variants.append(Variant(schedule, exprs))
 
         if not variants:
             return []
@@ -147,7 +148,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 +273,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 +853,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 +929,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 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

@@ -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,6 +2701,8 @@ def test_split_cond(self):
         eq2 = Eq(u.forward, u.forward + cos(time), implicit_dims=ct)
 
         op = Operator([eq0, eq1, eq2])
+        op(ime=5)
+
         cond = FindNodes(Conditional).visit(op)
         assert len(cond) == 3
         # The alias should have been lifted out of the condition
@@ -2721,6 +2724,8 @@ 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(ime=5)
+
         cond = FindNodes(Conditional).visit(op)
         assert len(cond) == 3
         # The alias should have been lifted out of the condition
@@ -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