compiler: add rudimentary support for multi-cond buffering

.github/workflows/docker-bases.yaml

@@ -40,7 +40,7 @@ jobs:
   ############## Basic gcc CPU ##########################
   #######################################################
   deploy-cpu-bases:
-    if: inputs.cpu
+    if: ${{ github.event_name != 'workflow_dispatch' || inputs.cpu }}
     name: "cpu-base-${{ matrix.arch }}-gcc${{ matrix.gcc }}"
     runs-on: ${{ matrix.runner }}
     env:
@@ -88,7 +88,7 @@ jobs:
           tags: "devitocodes/bases:cpu-gcc${{ matrix.gcc }}-${{ matrix.arch }}"
 
   deploy-cpu-bases-manifest:
-    if: inputs.cpu
+    if: ${{ github.event_name != 'workflow_dispatch' || inputs.cpu }}
     name: "cpu-base-manifest"
     runs-on: ubuntu-latest
     needs: deploy-cpu-bases
@@ -123,7 +123,7 @@ jobs:
   ############## Intel OneApi CPU #######################
   #######################################################
   deploy-oneapi-bases:
-    if: inputs.intel
+    if: ${{ github.event_name != 'workflow_dispatch' || inputs.intel }}
     name: "oneapi-base"
     runs-on: ubuntu-latest
     env:
@@ -181,7 +181,7 @@ jobs:
   ################### Nvidia nvhpc ######################
   #######################################################
   deploy-nvidia-bases:
-    if: inputs.nvidia
+    if: ${{ github.event_name != 'workflow_dispatch' || inputs.nvidia }}
     name: "nvidia-bases-${{ matrix.arch }}"
     runs-on: ${{ matrix.runner }}
     env:
@@ -270,7 +270,7 @@ jobs:
           tags: "devitocodes/bases:cpu-nvc${{ matrix.extra_tag }}-${{ matrix.arch }}"
 
   deploy-nvidia-bases-manifest:
-    if: inputs.nvidia
+    if: ${{ github.event_name != 'workflow_dispatch' || inputs.nvidia }}
     name: "nvidia-base-manifest"
     runs-on: ubuntu-latest
     needs: deploy-nvidia-bases
@@ -305,7 +305,7 @@ jobs:
   ##################### AMD #############################
   #######################################################
   deploy-amd-bases:
-    if: inputs.amd
+    if: ${{ github.event_name != 'workflow_dispatch' || inputs.amd }}
     name: "amd-base"
     runs-on: ["self-hosted", "amdgpu"]
     env:

devito/ir/clusters/algorithms.py

@@ -254,6 +254,7 @@ def guard(clusters):
 
             # Chain together all `cds` conditions from all expressions in `c`
             guards = {}
+            mode = sympy.Or
             for cd in cds:
                 # `BOTTOM` parent implies a guard that lives outside of
                 # any iteration space, which corresponds to the placeholder None
@@ -270,6 +271,7 @@ def guard(clusters):
 
                 # Pull `cd` from any expr
                 condition = guards.setdefault(k, [])
+                mode = mode and cd.relation
                 for e in exprs:
                     try:
                         condition.append(e.conditionals[cd])
@@ -284,7 +286,8 @@ def guard(clusters):
                     conditionals.pop(cd, None)
                     exprs[i] = e.func(*e.args, conditionals=conditionals)
 
-            guards = {d: sympy.And(*v, evaluate=False) for d, v in guards.items()}
+            # Combination mode is And by default and Or if all conditions are
+            guards = {d: mode(*v, evaluate=False) for d, v in guards.items()}
 
             # Construct a guarded Cluster
             processed.append(c.rebuild(exprs=exprs, guards=guards))

devito/ir/equations/equation.py

@@ -223,7 +223,7 @@ def __new__(cls, *args, **kwargs):
             else:
                 cond = diff2sympy(lower_exprs(d.condition))
                 if d._factor is not None:
-                    cond = sympy.And(cond, GuardFactor(d))
+                    cond = d.relation(cond, GuardFactor(d))
                 conditionals[d] = cond
             # Replace dimension with index
             index = d.index

devito/types/dimension.py

@@ -860,6 +860,8 @@ class ConditionalDimension(DerivedDimension):
         If True, use `self`, rather than the parent Dimension, to
         index into arrays. A typical use case is when arrays are accessed
         indirectly via the ``condition`` expression.
+    relation: Or/And, default=And
+        How this ConditionalDimension will be combined with other ones.
 
     Examples
     --------
@@ -913,10 +915,10 @@ class ConditionalDimension(DerivedDimension):
     is_Conditional = True
 
     __rkwargs__ = DerivedDimension.__rkwargs__ + \
-        ('factor', 'condition', 'indirect')
+        ('factor', 'condition', 'indirect', 'relation')
 
     def __init_finalize__(self, name, parent=None, factor=None, condition=None,
-                          indirect=False, **kwargs):
+                          indirect=False, relation=sympy.And, **kwargs):
         # `parent=None` degenerates to a ConditionalDimension outside of
         # any iteration space
         if parent is None:
@@ -937,6 +939,7 @@ def __init_finalize__(self, name, parent=None, factor=None, condition=None,
 
         self._condition = condition
         self._indirect = indirect
+        self._relation = relation
 
     @property
     def uses_symbolic_factor(self):
@@ -978,6 +981,10 @@ def condition(self):
     def indirect(self):
         return self._indirect
 
+    @property
+    def relation(self):
+        return self._relation
+
     @cached_property
     def free_symbols(self):
         retval = set(super().free_symbols)

docker/Dockerfile.intel

@@ -26,7 +26,7 @@ FROM base AS oneapi
 
 # Download the key to system keyring
 # https://www.intel.com/content/www/us/en/develop/documentation/installation-guide-for-intel-oneapi-toolkits-linux/top/installation/install-using-package-managers/apt.html#apt
-SHELL /bin/bash -o pipefail
+SHELL ["/bin/bash", "-o", "pipefail", "-c"]
 RUN wget --progress=dot:giga -O- https://apt.repos.intel.com/intel-gpg-keys/GPG-PUB-KEY-INTEL-SW-PRODUCTS.PUB | gpg --dearmor > /usr/share/keyrings/oneapi-archive-keyring.gpg
 RUN echo "deb [signed-by=/usr/share/keyrings/oneapi-archive-keyring.gpg] https://apt.repos.intel.com/oneapi all main" > /etc/apt/sources.list.d/oneAPI.list
 
@@ -37,7 +37,7 @@ RUN apt-get update -y && \
 
 # Drivers mandatory for intel gpu
 # https://dgpu-docs.intel.com/driver/installation.html#ubuntu-install-steps
-SHELL /bin/bash -o pipefail
+SHELL ["/bin/bash", "-o", "pipefail", "-c"]
 RUN wget -qO - https://repositories.intel.com/graphics/intel-graphics.key | gpg --dearmor > /usr/share/keyrings/intel-graphics.gpg
 RUN echo "deb [arch=amd64 signed-by=/usr/share/keyrings/intel-graphics.gpg] https://repositories.intel.com/graphics/ubuntu jammy unified" >  /etc/apt/sources.list.d/intel-gpu-jammy.list
 

docker/Dockerfile.nvidia

@@ -19,7 +19,7 @@ RUN apt-get update && \
                           dh-autoreconf python3-venv python3-dev python3-pip
 
 # nodesource: nvdashboard requires nodejs>=10
-SHELL /bin/bash -o pipefail
+SHELL ["/bin/bash", "-o", "pipefail", "-c"]
 RUN curl https://developer.download.nvidia.com/hpc-sdk/ubuntu/DEB-GPG-KEY-NVIDIA-HPC-SDK | gpg --yes --dearmor -o /usr/share/keyrings/nvidia-hpcsdk-archive-keyring.gpg
 RUN arch="$(uname -m)" && \
     case "$arch" in \
@@ -92,7 +92,7 @@ ENV UCX_TLS=cuda,cuda_copy,cuda_ipc,sm,shm,self
 #ENV UCX_TLS=cuda,cuda_copy,cuda_ipc,sm,shm,self,rc_x,gdr_copy
 
 # Make simlink for path setup since ENV doesn't accept shell commands.
-SHELL /bin/bash -o pipefail
+SHELL ["/bin/bash", "-o", "pipefail", "-c"]
 RUN arch="$(uname -m)" && \
     case "$arch" in \
         x86_64) linux=Linux_x86_64 ;; \

examples/userapi/05_conditional_dimension.ipynb

@@ -179,6 +179,8 @@
       "        If True, use `self`, rather than the parent Dimension, to\n",
       "        index into arrays. A typical use case is when arrays are accessed\n",
       "        indirectly via the ``condition`` expression.\n",
+      "    relation: Or/And, default=And\n",
+      "        How this ConditionalDimension will be combined with other ones.\n",
       "\n",
       "    Examples\n",
       "    --------\n",
@@ -248,9 +250,7 @@
   {
    "cell_type": "code",
    "execution_count": 5,
-   "metadata": {
-    "scrolled": true
-   },
+   "metadata": {},
    "outputs": [
     {
      "name": "stderr",
@@ -323,7 +323,7 @@
       "START(section0)\n",
       "for (int x = x_m; x <= x_M; x += 1)\n",
       "{\n",
-      "  #pragma omp simd aligned(f:32)\n",
+      "  #pragma omp simd aligned(f:64)\n",
       "  for (int y = y_m; y <= y_M; y += 1)\n",
       "  {\n",
       "    if (f[x + 1][y + 1] > 0)\n",
@@ -399,7 +399,7 @@
       "START(section0)\n",
       "for (int x = x_m; x <= x_M; x += 1)\n",
       "{\n",
-      "  #pragma omp simd aligned(f,g:32)\n",
+      "  #pragma omp simd aligned(f,g:64)\n",
       "  for (int y = y_m; y <= y_M; y += 1)\n",
       "  {\n",
       "    if (y < 5 && g[x + 1][y + 1] != 0)\n",
@@ -551,13 +551,6 @@
    "execution_count": 9,
    "metadata": {},
    "outputs": [
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "Operator `Kernel` ran in 0.01 s\n"
-     ]
-    },
     {
      "name": "stdout",
      "output_type": "stream",
@@ -570,7 +563,20 @@
       "    f[i / cif] = g[i];\n",
       "  }\n",
       "}\n",
-      "STOP(section0,timers)\n",
+      "STOP(section0,timers)\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Operator `Kernel` ran in 0.01 s\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
       "\n",
       " Data in g \n",
       " [ 0.  1.  2.  3.  4.  5.  6.  7.  8.  9. 10. 11. 12. 13. 14. 15.]\n",
@@ -831,7 +837,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.11.2"
+   "version": "3.11.0rc1"
   }
  },
  "nbformat": 4,

tests/test_buffering.py

@@ -1,9 +1,10 @@
 import numpy as np
 import pytest
+from sympy import Or
 
 from conftest import skipif
 from devito import (
-    ConditionalDimension, Constant, Eq, Grid, Operator, SubDimension, SubDomain,
+    CondEq, ConditionalDimension, Constant, Eq, Grid, Operator, SubDimension, SubDomain,
     TimeFunction, configuration, switchconfig
 )
 from devito.arch.archinfo import AppleArm
@@ -751,3 +752,36 @@ def test_buffer_reuse():
 
     assert all(np.all(usave.data[i-1] == i) for i in range(1, nt + 1))
     assert all(np.all(vsave.data[i-1] == i + 1) for i in range(1, nt + 1))
+
+
+def test_multi_cond():
+    grid = Grid((3, 3))
+    nt = 5
+
+    x, y = grid.dimensions
+
+    factor = 2
+    ntmod = (nt - 1) * factor + 1
+
+    ct1 = ConditionalDimension(name="ct1", parent=grid.time_dim,
+                               factor=factor, relation=Or)
+    ctend = ConditionalDimension(name="ctend", parent=grid.time_dim,
+                                 condition=CondEq(grid.time_dim, ntmod - 2),
+                                 relation=Or)
+
+    f = TimeFunction(grid=grid, name='f', time_order=0,
+                     space_order=0, save=nt, time_dim=ct1)
+    T = TimeFunction(grid=grid, name='T', time_order=0, space_order=0)
+
+    eqs = [Eq(T, grid.time_dim)]
+    # this to save times from 0 to nt - 2
+    eqs.append(Eq(f, T))
+    # this to save the last time sample nt - 1
+    eqs.append(Eq(f.forward, T+1, implicit_dims=ctend))
+
+    # run operator with buffering
+    op = Operator(eqs, opt=('streaming', 'buffering'))
+    op.apply(time_m=0, time_M=ntmod-2)
+
+    for i in range(nt):
+        assert np.allclose(f.data[i], i*2)