comiler: better handle buffer ispace

Author: mloubout

.github/workflows/docker-bases.yml

@@ -18,16 +18,16 @@ on:
     inputs:
       cpu:
         type: boolean
-        default: false
+        default: true
       nvidia:
         type: boolean
-        default: false
+        default: true
       amd:
         type: boolean
-        default: false
+        default: true
       intel:
         type: boolean
-        default: false
+        default: true
 
       tags:
         description: "Build compiler bases"

devito/passes/clusters/buffering.py

@@ -461,12 +461,13 @@ def itdims(self):
     @cached_property
     def ispace(self):
         # The IterationSpace within which the buffer will be accessed
+
         # NOTE: The `key` is to avoid Clusters including `f` but not directly
         # using it in an expression, such as HaloTouch Clusters
         def key(c):
             bufferdim = any(i in c.ispace.dimensions for i in self.bdims)
-            timeonly = all(d.is_Time for d in c.ispace.dimensions)
-            return bufferdim or timeonly
+            xd_only = all(d._defines & self.xd._defines for d in c.ispace.dimensions)
+            return bufferdim or xd_only
 
         ispaces = set()
         for c in self.clusters:
@@ -478,14 +479,13 @@ def key(c):
                 continue
 
             # Iterations space and buffering dims
-            ispace = c.ispace
-            edims = [d for d in self.bdims if d not in ispace.dimensions]
+            edims = [d for d in self.bdims if d not in c.ispace.dimensions]
             if not edims:
-                ispaces.add(ispace)
+                ispaces.add(c.ispace)
             else:
                 # Add all missing buffering dimensions and reorder to
                 # avoid duplicates with different ordering
-                ispaces.add(ispace.insert(self.dim, edims).reorder())
+                ispaces.add(c.ispace.insert(self.dim, edims).reorder())
 
         if len(ispaces) > 1:
             # Best effort to make buffering work in the presence of multiple

docker/Dockerfile.nvidia

@@ -20,20 +20,32 @@ RUN apt-get update && \
 
 # nodesource: nvdashboard requires nodejs>=10
 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 echo 'deb [trusted=yes, signed-by=/usr/share/keyrings/nvidia-hpcsdk-archive-keyring.gpg] https://developer.download.nvidia.com/hpc-sdk/ubuntu/amd64 /' | tee /etc/apt/sources.list.d/nvhpc.list
+RUN arch="$(uname -m)" && \
+    case "$arch" in \
+        x86_64) nvplat=amd64 ;; \
+        aarch64|arm64) nvplat=arm64 ;; \
+        *) echo "Unsupported architecture: $arch" >&2; exit 1 ;; \
+    esac && \
+    echo "deb [trusted=yes, signed-by=/usr/share/keyrings/nvidia-hpcsdk-archive-keyring.gpg] https://developer.download.nvidia.com/hpc-sdk/ubuntu/${nvplat} /" | tee /etc/apt/sources.list.d/nvhpc.list
 RUN apt-key update *&& apt-get update -y
 
 # Install nvhpc. `nvhpc` is the alias for the latest avaialble version
 ARG ver=nvhpc
 # We use the standard apt-get for the default latest nvhpc. For earlier version, apt has a bug that it will always
 # install the latest nvhpc-x-y no matter which version nvhpc-x-z is requested which would double (extra 10Gb) the size of the image.
 # So for specific version we directly download the specific deb and install it.
-RUN if [ "$ver" = "nvhpc" ]; then \
+RUN arch="$(uname -m)" && \
+    case "$arch" in \
+        x86_64) nvplat=amd64 ;; \
+        aarch64|arm64) nvplat=arm64 ;; \
+        *) echo "Unsupported architecture: $arch" >&2; exit 1 ;; \
+    esac && \
+    if [ "$ver" = "nvhpc" ]; then \
         apt-get install -y -q --allow-unauthenticated ${ver}; \
     else \
         export year=$(echo $ver | cut -d "-" -f 2) && export minor=$(echo $ver | cut -d "-" -f 3) && \
-        wget https://developer.download.nvidia.com/hpc-sdk/ubuntu/amd64/nvhpc_${year}.${minor}_amd64.deb && \
-        apt-get install --allow-unauthenticated -y -q ./nvhpc_${year}.${minor}_amd64.deb; \
+        wget https://developer.download.nvidia.com/hpc-sdk/ubuntu/${nvplat}/nvhpc_${year}.${minor}_${nvplat}.deb && \
+        apt-get install --allow-unauthenticated -y -q ./nvhpc_${year}.${minor}_${nvplat}.deb; \
     fi;
 
 # Nodejs https://github.com/nodesource/distributions