comiler: better handle buffer ispace

mloubout · mloubout · commit edeb3b6a1659 · 2025-10-29T12:47:27.000-04:00
diff --git a/.github/workflows/docker-bases.yml b/.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"
@@ -227,7 +227,9 @@ jobs:
           push: true
           target: "nvc"
           platforms: ${{ matrix.platform }}
-          build-args: "arch=nvc"
+          build-args: |
+            arch=nvc
+            ver=nvhpc-25-7
           # Label (not tag) with runner name for traceability without changing image tags
           labels: builder-runner=${{ runner.name }}
           tags: "devitocodes/bases:nvidia-nvc-${{ matrix.arch }}"
@@ -240,7 +242,9 @@ jobs:
           push: true
           target: "nvcc"
           platforms: ${{ matrix.platform }}
-          build-args: "arch=nvcc"
+          build-args: |
+            arch=nvcc
+            ver=nvhpc-25-7
           labels: builder-runner=${{ runner.name }}
           tags: "devitocodes/bases:nvidia-nvcc-${{ matrix.arch }}"
 
@@ -252,7 +256,9 @@ jobs:
           push: true
           target: "nvc-host"
           platforms: ${{ matrix.platform }}
-          build-args: "arch=nvc-host"
+          build-args: |
+            arch=nvc-host
+            ver=nvhpc-25-7
           labels: builder-runner=${{ runner.name }}
           tags: "devitocodes/bases:cpu-nvc-${{ matrix.arch }}"
 
diff --git a/.github/workflows/docker-devito.yml b/.github/workflows/docker-devito.yml
@@ -215,7 +215,7 @@ jobs:
             echo "No tags generated, skipping"
             exit 0
           fi
-          for tag in $TAGS; do
+          echo "$TAGS" | while read -r tag; do
             refs=""
             for arch in $ARCHES; do
               refs="$refs devitocodes/devito:${tag}-${arch}"
diff --git a/devito/passes/clusters/buffering.py b/devito/passes/clusters/buffering.py
@@ -405,8 +405,10 @@ def generate_buffers(clusters, key, sregistry, options, **kwargs):
         name = sregistry.make_name(prefix='%sb' % f.name)
         # We specify the padding to match the input Function's one, so that
         # the array can be used in place of the Function with valid strides
+        # Plain Array do not track mapped so we default to no padding
+        padding = 0 if cls is Array else f.padding
         mapper[f] = cls(name=name, dimensions=dimensions, dtype=f.dtype,
-                        padding=f.padding, grid=f.grid, halo=f.halo,
+                        padding=padding, grid=f.grid, halo=f.halo,
                         space='mapped', mapped=f, f=f)
 
     return mapper
@@ -461,12 +463,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 +481,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
diff --git a/devito/passes/iet/languages/CXX.py b/devito/passes/iet/languages/CXX.py
@@ -67,6 +67,10 @@ def std_arith(prefix=None):
 
 
 def split_pointer(i, idx):
+    """
+    Splits complex pointer std::complex<T> *a as
+    (float *)(&a)[idx] for real/imag parts.
+    """
     dtype = i.dtype(0).real.__class__
     ptr = cast(dtype, stars='*')(Byref(i), reinterpret=True)
     return IndexedPointer(ptr, idx)
diff --git a/devito/passes/iet/languages/utils.py b/devito/passes/iet/languages/utils.py
@@ -18,24 +18,24 @@ def _atomic_add_split(i, pragmas, real, imag):
     lhs, rhs = i.expr.lhs, i.expr.rhs
     if (np.issubdtype(lhs.dtype, np.complexfloating)
        and np.issubdtype(rhs.dtype, np.complexfloating)):
-        # Complex i, complex j
+        # Complex lhs, complex rhs
         # Atomic add real and imaginary parts separately
         lhsr, rhsr = real(lhs), Real(rhs)
         lhsi, rhsi = imag(lhs), Imag(rhs)
-        real = i._rebuild(expr=i.expr._rebuild(lhs=lhsr, rhs=rhsr),
-                          pragmas=pragmas)
-        imag = i._rebuild(expr=i.expr._rebuild(lhs=lhsi, rhs=rhsi),
-                          pragmas=pragmas)
-        return List(body=[real, imag])
+        real_eq = i._rebuild(expr=i.expr._rebuild(lhs=lhsr, rhs=rhsr),
+                             pragmas=pragmas)
+        imag_eq = i._rebuild(expr=i.expr._rebuild(lhs=lhsi, rhs=rhsi),
+                             pragmas=pragmas)
+        return List(body=[real_eq, imag_eq])
 
     elif (np.issubdtype(lhs.dtype, np.complexfloating)
           and not np.issubdtype(rhs.dtype, np.complexfloating)):
-        # Complex i, real j
-        # Atomic add j to real part of i
+        # Complex lhs, real rhs
+        # Atomic add rhs to real part of lhs
         lhsr, rhsr = real(lhs), rhs
-        real = i._rebuild(expr=i.expr._rebuild(lhs=lhsr, rhs=rhsr),
-                          pragmas=pragmas)
-        return real
+        real_eq = i._rebuild(expr=i.expr._rebuild(lhs=lhsr, rhs=rhsr),
+                             pragmas=pragmas)
+        return real_eq
     else:
         # Real i, complex j
         raise InvalidOperator("Atomic add not implemented for real "
diff --git a/docker/Dockerfile.nvidia b/docker/Dockerfile.nvidia
@@ -11,7 +11,7 @@ FROM ubuntu:22.04 AS sdk-base
 
 SHELL ["/bin/bash", "-c"]
 
-ENV DEBIAN_FRONTEND noninteractive
+ENV DEBIAN_FRONTEND=noninteractive
 
 # Install python
 RUN apt-get update && \
@@ -20,20 +20,34 @@ 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; \
+        export year=$(echo $ver | cut -d "-" -f 2) && \
+        export minor=$(echo $ver | cut -d "-" -f 3) && \
+        wget -O nvhpc.deb "https://developer.download.nvidia.com/hpc-sdk/ubuntu/${nvplat}/nvhpc_${year}.${minor}_${nvplat}.deb" \
+        || wget -O nvhpc.deb "https://developer.download.nvidia.com/hpc-sdk/ubuntu/${nvplat}/nvhpc_${year}.${minor}-0_${nvplat}.deb" && \
+        apt-get install --allow-unauthenticated -y -q ./nvhpc.deb; \
     fi;
 
 # Nodejs https://github.com/nodesource/distributions
@@ -43,12 +57,15 @@ RUN apt-get update && apt-get install -y -q \
         liblapack-dev libblas-dev \
         libibverbs-dev libmlx4-1 libmlx5-1 ibutils \
         # Devito Jupyter Notebooks and Ux experience
-        nodejs ffmpeg gcc-offload-nvptx \
+        nodejs ffmpeg \
         texlive-latex-extra texlive-fonts-recommended dvipng cm-super
 
+# nvptx only available on x86_64
+RUN apt-get install -y -q gcc-offload-nvptx || echo "Skipping nvptx compiler installation on non-x86_64 architecture"
+
 # nvidia-container-runtime
-ENV NVIDIA_VISIBLE_DEVICES all
-ENV NVIDIA_DRIVER_CAPABILITIES compute,utility
+ENV NVIDIA_VISIBLE_DEVICES=all
+ENV NVIDIA_DRIVER_CAPABILITIES=compute,utility
 
 # MPI ROOT USER DEFAULTS
 ENV OMPI_ALLOW_RUN_AS_ROOT=1
@@ -72,13 +89,19 @@ 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.
-RUN export NVARCH=$(ls -1 /opt/nvidia/hpc_sdk/Linux_x86_64/ | grep '\.' | head -n 1) && \
-    export CUDA_V=$(ls /opt/nvidia/hpc_sdk/Linux_x86_64/${NVARCH}/cuda/ | grep '\.') && \
-    ln -sf /opt/nvidia/hpc_sdk/Linux_x86_64/${NVARCH} /opt/nvhpc && \
-    ln -sf /opt/nvidia/hpc_sdk/Linux_x86_64/${NVARCH}/cuda/${CUDA_V}/extras/CUPTI /opt/CUPTI && \
-    ln -sf /opt/nvidia/hpc_sdk/Linux_x86_64/comm_libs/${CUDA_V}/nvshmem /opt/nvhpc/comm_libs/nvshmem && \
-    ln -sf /opt/nvidia/hpc_sdk/Linux_x86_64/comm_libs/${CUDA_V}/nccl /opt/nvhpc/comm_libs/nccl && \
-    ln -sf /opt/nvidia/hpc_sdk/Linux_x86_64/${NVARCH}/cuda/${CUDA_V}/compute-sanitizer/compute-sanitizer /opt/nvhpc/compilers/bin/compute-sanitizer
+RUN arch="$(uname -m)" && \
+    case "$arch" in \
+        x86_64) linux=Linux_x86_64 ;; \
+        aarch64|arm64) linux=Linux_aarch64 ;; \
+        *) echo "Unsupported architecture: $arch" >&2; exit 1 ;; \
+    esac && \
+    export NVARCH=$(ls -1 /opt/nvidia/hpc_sdk/${linux}/ | grep '\.' | head -n 1) && \
+    export CUDA_V=$(ls /opt/nvidia/hpc_sdk/${linux}/${NVARCH}/cuda/ | grep '\.') && \
+    ln -sf /opt/nvidia/hpc_sdk/${linux}/${NVARCH} /opt/nvhpc && \
+    ln -sf /opt/nvidia/hpc_sdk/${linux}/${NVARCH}/cuda/${CUDA_V}/extras/CUPTI /opt/CUPTI && \
+    ln -sf /opt/nvidia/hpc_sdk/${linux}/comm_libs/${CUDA_V}/nvshmem /opt/nvhpc/comm_libs/nvshmem && \
+    ln -sf /opt/nvidia/hpc_sdk/${linux}/comm_libs/${CUDA_V}/nccl /opt/nvhpc/comm_libs/nccl && \
+    ln -sf /opt/nvidia/hpc_sdk/${linux}/${NVARCH}/cuda/${CUDA_V}/compute-sanitizer/compute-sanitizer /opt/nvhpc/compilers/bin/compute-sanitizer
 
 # Starting nvhpc 23.5 and cuda 12.1, hpcx and openmpi are inside the cuda version folder, only the bin is in the comm_libs path
 RUN export CUDA_V=$(/opt/nvhpc/cuda/bin/nvcc --version | sed -n 's/^.*release \([0-9]\+\.[0-9]\+\).*$/\1/p') && \
@@ -102,12 +125,12 @@ RUN echo "$HPCSDK_HOME/cuda/lib" >> /etc/ld.so.conf.d/nvidia.conf && \
     
 # Compiler, CUDA, and Library paths
 # CUDA_HOME has been deprecated but keep for now because of other dependencies (@mloubout).
-ENV CUDA_HOME $HPCSDK_HOME/cuda
-ENV NVHPC_CUDA_HOME $HPCSDK_HOME/cuda
-ENV CUDA_ROOT $HPCSDK_HOME/cuda/bin
-ENV PATH $HPCSDK_HOME/compilers/bin:$HPCSDK_HOME/cuda/bin:$HPCSDK_HOME/comm_libs/mpi/bin:${PATH}
-ENV LD_LIBRARY_PATH $HPCSDK_HOME/cuda/lib:$HPCSDK_HOME/cuda/lib64:$HPCSDK_HOME/compilers/lib:$HPCSDK_HOME/math_libs/lib64:$HPCSDK_HOME/comm_libs/mpi/lib:$HPCSDK_CUPTI/lib64:bitcomp_DIR:${LD_LIBRARY_PATH}
-ENV CPATH $HPCSDK_HOME/comm_libs/mpi/include:$HPCSDK_HOME/comm_libs/nvshmem/include:$HPCSDK_HOME/comm_libs/nccl/include:$HPCSDK_HOME/math_libs/include:${CPATH}
+ENV CUDA_HOME=$HPCSDK_HOME/cuda
+ENV NVHPC_CUDA_HOME=$HPCSDK_HOME/cuda
+ENV CUDA_ROOT=$HPCSDK_HOME/cuda/bin
+ENV PATH=$HPCSDK_HOME/compilers/bin:$HPCSDK_HOME/cuda/bin:$HPCSDK_HOME/comm_libs/mpi/bin:${PATH}
+ENV LD_LIBRARY_PATH=$HPCSDK_HOME/cuda/lib:$HPCSDK_HOME/cuda/lib64:$HPCSDK_HOME/compilers/lib:$HPCSDK_HOME/math_libs/lib64:$HPCSDK_HOME/comm_libs/mpi/lib:$HPCSDK_CUPTI/lib64:bitcomp_DIR
+ENV CPATH=$HPCSDK_HOME/comm_libs/mpi/include:$HPCSDK_HOME/comm_libs/nvshmem/include:$HPCSDK_HOME/comm_libs/nccl/include:$HPCSDK_HOME/math_libs/include
 
 # MPI
 RUN rm -f  $HPCSDK_HOME/comm_libs/mpi && \
