compiler: prevent undefined Temp through global init

Author: mloubout

devito/passes/clusters/aliases.py

@@ -25,6 +25,7 @@
     CustomDimension, Eq, Hyperplane, IncrDimension, Indexed, ModuloDimension, Size,
     StencilDimension, Symbol, Temp, TempArray, TempFunction
 )
+from devito.types.dimension import ConditionalDimension, SubsamplingFactor
 from devito.types.grid import MultiSubDimension
 
 __all__ = ['cire']
@@ -923,15 +924,25 @@ def lower_schedule(schedule, meta, sregistry, opt_ftemps, opt_min_dtype,
             callback = lambda idx: obj[  # noqa: B023
                 [i + s for i, s in zip(idx, shift, strict=True)]  # noqa: B023
             ]
+            guards = meta.guards
         else:
             # Degenerate case: scalar expression
             assert writeto.size == 0
 
-            dtype = sympy_dtype(pivot, base=meta.dtype, smin=opt_min_dtype)
-            obj = Temp(name=name, dtype=dtype)
-            expression = Eq(obj, uxreplace(pivot, subs))
+            guards = None
+            is_cond = any(isinstance(d, (SubsamplingFactor, ConditionalDimension))
+                          for d in pivot.free_symbols)
+            if meta.guards and is_cond:
+                # 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)
+                expression = Eq(obj, uxreplace(pivot, subs))
 
-            callback = lambda idx: obj  # noqa: B023
+                callback = lambda idx: obj  # noqa: B023
 
         # Create the substitution rules for the aliasing expressions
         subs.update({
@@ -957,7 +968,8 @@ def lower_schedule(schedule, meta, sregistry, opt_ftemps, opt_min_dtype,
             properties[Hyperplane(writeto.itdims)] = {SEPARABLE}
 
         # Finally, build the alias Cluster
-        clusters.append(Cluster(expression, ispace, meta.guards, properties))
+        if expression is not None:
+            clusters.append(Cluster(expression, ispace, guards, properties))
 
     return clusters, subs
 

devito/passes/clusters/cse.py

@@ -40,6 +40,12 @@ def cse_dtype(exprdtype, cdtype):
     Return the dtype of a CSE temporary given the dtype of the expression to be
     captured and the cluster's dtype.
     """
+    if np.issubdtype(cdtype, np.floating) and np.issubdtype(exprdtype, np.integer):
+        # Integer expression and floating-point cluster: promote to the floating point
+        # np.promote_types upcast integers (e.g int32 -> Float64) so we
+        # need to ensure that the promoted type is not larger than the cluster's dtype
+        return cdtype
+
     if np.issubdtype(cdtype, np.complexfloating):
         return np.promote_types(exprdtype, cdtype(0).real.__class__).type
     else:
@@ -97,8 +103,9 @@ def cse(cluster, sregistry=None, options=None, **kwargs):
     if cluster.is_fence:
         return cluster
 
-    make_dtype = lambda e: cse_dtype(e.dtype, dtype)
-    make = lambda e: CTemp(name=sregistry.make_name(), dtype=make_dtype(e))
+    def make(e):
+        edtype = cse_dtype(e.dtype, dtype)
+        return CTemp(name=sregistry.make_name(), dtype=edtype)
 
     exprs = _cse(cluster, make, min_cost=min_cost, mode=mode)
 

devito/types/basic.py

@@ -1589,7 +1589,7 @@ def _rebuild(self, *args, **kwargs):
         comps = [f.func(*args, name=f.name.replace(self.name, newname), **kwargs)
                  for f in self.flat()]
         # Rebuild the matrix with the new components
-        return self._new(comps)
+        return self._new(*self.shape, comps)
 
     func = _rebuild
 

examples/mpi/overview.ipynb

@@ -460,7 +460,7 @@
        "  _MM_SET_DENORMALS_ZERO_MODE(_MM_DENORMALS_ZERO_ON);\n",
        "  _MM_SET_FLUSH_ZERO_MODE(_MM_FLUSH_ZERO_ON);\n",
        "\n",
-       "  float r0 = 1.0F/h_x;\n",
+       "  const float r0 = 1.0F/h_x;\n",
        "\n",
        "  for (int time = time_m, t0 = (time)%(2), t1 = (time + 1)%(2); time <= time_M; time += 1, t0 = (time)%(2), t1 = (time + 1)%(2))\n",
        "  {\n",

examples/performance/00_overview.ipynb

@@ -508,7 +508,7 @@
       "}\n",
       "STOP(section0,timers)\n",
       "\n",
-      "float r1 = 1.0F/h_y;\n",
+      "const float r1 = 1.0F/h_y;\n",
       "\n",
       "for (int time = time_m, t0 = (time)%(2), t1 = (time + 1)%(2); time <= time_M; time += 1, t0 = (time)%(2), t1 = (time + 1)%(2))\n",
       "{\n",
@@ -572,7 +572,13 @@
       "+            u[t1][x + 4][y + 4][z + 4] = (f[x + 1][y + 1][z + 1]*f[x + 1][y + 1][z + 1])*((-6.66666667e-1F*r0)*(8.33333333e-2F*r0*u[t0][x + 4][y + 1][z + 4] - 6.66666667e-1F*r0*u[t0][x + 4][y + 2][z + 4] + 6.66666667e-1F*r0*u[t0][x + 4][y + 4][z + 4] - 8.33333333e-2F*r0*u[t0][x + 4][y + 5][z + 4]) + (-8.33333333e-2F*r0)*(8.33333333e-2F*r0*u[t0][x + 4][y + 4][z + 4] - 6.66666667e-1F*r0*u[t0][x + 4][y + 5][z + 4] + 6.66666667e-1F*r0*u[t0][x + 4][y + 7][z + 4] - 8.33333333e-2F*r0*u[t0][x + 4][y + 8][z + 4]) + (8.33333333e-2F*r0)*(8.33333333e-2F*r0*u[t0][x + 4][y][z + 4] - 6.66666667e-1F*r0*u[t0][x + 4][y + 1][z + 4] + 6.66666667e-1F*r0*u[t0][x + 4][y + 3][z + 4] - 8.33333333e-2F*r0*u[t0][x + 4][y + 4][z + 4]) + (6.66666667e-1F*r0)*(8.33333333e-2F*r0*u[t0][x + 4][y + 3][z + 4] - 6.66666667e-1F*r0*u[t0][x + 4][y + 4][z + 4] + 6.66666667e-1F*r0*u[t0][x + 4][y + 6][z + 4] - 8.33333333e-2F*r0*u[t0][x + 4][y + 7][z + 4]))*sinf(f[x + 1][y + 1][z + 1]);\n",
       "           }\n",
       "         }\n",
-      "       }\n",
+      "       }\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
       "\n"
      ]
     }
@@ -652,7 +658,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 13,
+   "execution_count": 12,
    "metadata": {},
    "outputs": [
     {
@@ -712,7 +718,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 14,
+   "execution_count": 13,
    "metadata": {},
    "outputs": [
     {
@@ -753,7 +759,14 @@
       "    }\n",
       "  }\n",
       "  STOP(section0,timers)\n",
-      "}\n"
+      "}"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n"
      ]
     }
    ],
@@ -772,7 +785,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 15,
+   "execution_count": 14,
    "metadata": {},
    "outputs": [
     {
@@ -863,7 +876,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 16,
+   "execution_count": 15,
    "metadata": {},
    "outputs": [
     {
@@ -919,7 +932,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 17,
+   "execution_count": 16,
    "metadata": {},
    "outputs": [
     {
@@ -976,7 +989,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 18,
+   "execution_count": 17,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -994,7 +1007,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 19,
+   "execution_count": 18,
    "metadata": {},
    "outputs": [
     {
@@ -1044,7 +1057,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 20,
+   "execution_count": 19,
    "metadata": {},
    "outputs": [
     {
@@ -1112,7 +1125,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 21,
+   "execution_count": 20,
    "metadata": {},
    "outputs": [
     {
@@ -1192,7 +1205,7 @@
       "  }\n",
       "  STOP(section0,timers)\n",
       "\n",
-      "  float r1 = 1.0F/h_y;\n",
+      "  const float r1 = 1.0F/h_y;\n",
       "\n",
       "  for (int time = time_m, t0 = (time)%(2), t1 = (time + 1)%(2); time <= time_M; time += 1, t0 = (time)%(2), t1 = (time + 1)%(2))\n",
       "  {\n",
@@ -1279,7 +1292,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 22,
+   "execution_count": 21,
    "metadata": {},
    "outputs": [
     {
@@ -1304,7 +1317,7 @@
       "}\n",
       "STOP(section0,timers)\n",
       "\n",
-      "float r1 = 1.0F/h_y;\n",
+      "const float r1 = 1.0F/h_y;\n",
       "\n",
       "for (int time = time_m, t0 = (time)%(2), t1 = (time + 1)%(2); time <= time_M; time += 1, t0 = (time)%(2), t1 = (time + 1)%(2))\n",
       "{\n",
@@ -1369,7 +1382,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 23,
+   "execution_count": 22,
    "metadata": {},
    "outputs": [
     {
@@ -1404,7 +1417,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 24,
+   "execution_count": 23,
    "metadata": {},
    "outputs": [
     {
@@ -1483,7 +1496,7 @@
       "  }\n",
       "  STOP(section0,timers)\n",
       "\n",
-      "  float r1 = 1.0F/h_y;\n",
+      "  const float r1 = 1.0F/h_y;\n",
       "\n",
       "  for (int time = time_m, t0 = (time)%(2), t1 = (time + 1)%(2); time <= time_M; time += 1, t0 = (time)%(2), t1 = (time + 1)%(2))\n",
       "  {\n",
@@ -1546,7 +1559,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 25,
+   "execution_count": 24,
    "metadata": {},
    "outputs": [
     {
@@ -1624,8 +1637,8 @@
       "  }\n",
       "  STOP(section0,timers)\n",
       "\n",
-      "  float r1 = 1.0F/h_x;\n",
-      "  float r2 = 1.0F/h_y;\n",
+      "  const float r1 = 1.0F/h_x;\n",
+      "  const float r2 = 1.0F/h_y;\n",
       "\n",
       "  for (int time = time_m, t0 = (time)%(2), t1 = (time + 1)%(2); time <= time_M; time += 1, t0 = (time)%(2), t1 = (time + 1)%(2))\n",
       "  {\n",
@@ -1718,7 +1731,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.10.12"
+   "version": "3.13.11"
   }
  },
  "nbformat": 4,

tests/test_dse.py

@@ -28,6 +28,7 @@
 )
 from devito.tools import as_tuple
 from devito.types import PrecomputedSparseTimeFunction, Scalar, Symbol
+from devito.types.misc import Temp
 from examples.seismic import AcquisitionGeometry, demo_model
 from examples.seismic.acoustic import AcousticWaveSolver
 from examples.seismic.tti import AnisotropicWaveSolver
@@ -2685,6 +2686,94 @@ def test_space_and_time_invariant_together(self):
             'tx0_blk0y0_blk0xyzyz'
         )
 
+    def test_split_cond(self):
+        grid = Grid((11, 11))
+        time = grid.time_dim
+
+        u = TimeFunction(name='u', grid=grid, time_order=2, space_order=2)
+
+        ct = ConditionalDimension(name='ct', parent=time, factor=2)
+        ct2 = ConditionalDimension(name='ct2', parent=time, factor=4)
+
+        eq0 = Eq(u.forward, u + cos(time), implicit_dims=ct)
+        eq1 = Eq(u.forward, u.forward + 1, implicit_dims=ct2)
+        eq2 = Eq(u.forward, u.forward + cos(time), implicit_dims=ct)
+
+        op = Operator([eq0, eq1, eq2])
+        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])
+        scalars = [i for i in FindSymbols().visit(op) if isinstance(i, Temp)]
+        assert len(scalars) == 1
+
+    def test_split_cond_multi_alias(self):
+        grid = Grid((11, 11))
+        time = grid.time_dim
+
+        u = TimeFunction(name='u', grid=grid, time_order=2, space_order=2)
+
+        ct = ConditionalDimension(name='ct', parent=time, factor=2)
+        ct2 = ConditionalDimension(name='ct2', parent=time, factor=4)
+
+        eq0 = Eq(u.forward, u + cos(time) + sin(time), implicit_dims=ct)
+        eq1 = Eq(u.forward, u.forward + 1, implicit_dims=ct2)
+        eq2 = Eq(u.forward, u.forward + cos(time) - sin(time), implicit_dims=ct)
+
+        op = Operator([eq0, eq1, eq2])
+        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])
+        scalars = [i for i in FindSymbols().visit(op) if isinstance(i, Temp)]
+        assert len(scalars) == 5
+
+    def test_multi_cond_no_split(self):
+        grid = Grid((11, 11))
+        time = grid.time_dim
+        u = TimeFunction(name='u', grid=grid, time_order=2, space_order=2)
+
+        ct = ConditionalDimension(name='ct', parent=time, factor=2)
+        ct2 = ConditionalDimension(name='ct2', parent=time, factor=4)
+
+        eq0 = Eq(u.forward, u + cos(time), implicit_dims=ct)
+        # Not hoisting the inite would have this equation split the time
+        # loop to initialize the alias for sin(time)
+        eq1 = Eq(u.forward, u.forward + sin(time), implicit_dims=ct2)
+        eq2 = Eq(u.forward, u.forward - sin(time), implicit_dims=ct)
+
+        op = Operator([eq0, eq1, eq2])
+
+        assert_structure(
+            op,
+            ['t', 't,x,y', 't,x,y', 't,x,y'],
+            'txyxyxy'
+        )
+
+        scalars = [i for i in FindSymbols().visit(op) if isinstance(i, Temp)]
+        assert len(scalars) == 4
+
+    def test_alias_with_conditional(self):
+        grid = Grid((11, 11))
+        time = grid.time_dim
+
+        u = TimeFunction(name='u', grid=grid, time_order=2, space_order=2)
+
+        ct = ConditionalDimension(name='ct', parent=time, factor=2)
+        ct2 = ConditionalDimension(name='ct2', parent=time, factor=4)
+
+        eq0 = Eq(u.forward, u + cos(ct), implicit_dims=ct)
+        eq1 = Eq(u.forward, u.forward + 1, implicit_dims=ct2)
+        eq2 = Eq(u.forward, u.forward + cos(ct), implicit_dims=ct)
+
+        op = Operator([eq0, eq1, eq2])
+        cond = FindNodes(Conditional).visit(op)
+        assert len(cond) == 3
+
+        # The alias should have been lifted out of the condition
+        scalars = [i for i in FindSymbols().visit(op) if isinstance(i, Temp)]
+        assert not scalars
+
 
 class TestIsoAcoustic: