ci: use GPU-aware MPI halos on GPU test runs (opt-out input)

.github/actions/run-mpas/action.yml

@@ -46,6 +46,13 @@ inputs:
       regardless of mpi-impl.
     required: false
     default: ''
+  gpu-aware-mpi:
+    description: >
+      Use GPU-aware MPI halo exchanges: sets config_gpu_aware_mpi = .true. in the namelist and the MPI library's
+      GPU-support environment (MPICH_GPU_SUPPORT_ENABLED / MPIR_CVAR_ENABLE_GPU for MPICH; nothing extra for OpenMPI, which is CUDA-aware by default). GPU (OpenACC) builds only, and only in
+      containers whose MPI is built with CUDA support. Default false; the GPU test workflows pass true.
+    required: false
+    default: 'false'
 
 outputs:
   log-dir:
@@ -114,6 +121,12 @@ runs:
         echo "=== Namelist ==="
         grep config_run_duration namelist.atmosphere
 
+    - name: Enable GPU-aware MPI halos
+      if: ${{ inputs.gpu-aware-mpi == 'true' }}
+      shell: bash
+      working-directory: ${{ steps.config.outputs.workdir }}
+      run: ${GITHUB_WORKSPACE}/.github/scripts/enable-gpu-aware-mpi.sh namelist.atmosphere
+
     - name: Run MPAS-A
       id: run
       shell: bash
@@ -149,13 +162,18 @@ runs:
           MPI_FLAGS="${MPI_FLAGS} ${{ inputs.mpi-run-flags }}"
         fi
 
+        if [ "${{ inputs.gpu-aware-mpi }}" = "true" ]; then
+          source "${GITHUB_WORKSPACE}/.github/scripts/gpu-aware-mpi-env.sh"
+        fi
+
         ulimit -s unlimited 2>/dev/null || echo "Warning: Could not set unlimited stack size"
 
         echo "=== Run configuration ==="
         echo "  Resolution: ${{ steps.config.outputs.resolution }}"
         echo "  Processors: ${{ inputs.num-procs }}"
         echo "  MPI_IMPL: ${MPI_IMPL:-auto}"
         echo "  MPI_FLAGS: ${MPI_FLAGS}"
+        echo "  GPU-aware MPI: ${{ inputs.gpu-aware-mpi }}"
         echo "  Stack limit: $(ulimit -s)"
         echo "  Available CPUs: $(nproc 2>/dev/null || echo unknown)"
         echo "  Available RAM: $(free -m 2>/dev/null | awk '/^Mem:/{print $2 "MB"}' || echo unknown)"

.github/actions/run-perturb-mpas/action.yml

@@ -33,6 +33,12 @@ inputs:
     description: 'MPI implementation (openmpi, mpich)'
     required: false
     default: 'openmpi'
+  gpu-aware-mpi:
+    description: >
+      Use GPU-aware MPI halo exchanges (config_gpu_aware_mpi = .true. in each member's namelist plus the MPI
+      library's GPU-support environment). GPU builds in containers with CUDA-enabled MPI only. Default false.
+    required: false
+    default: 'false'
   output-dir:
     description: 'Directory to collect trimmed history files'
     required: false
@@ -112,6 +118,10 @@ runs:
         if [ "${MPI_IMPL}" = "openmpi" ]; then
           MPI_FLAGS="${OPENMPI_RUN_FLAGS:---allow-run-as-root --oversubscribe}"
         fi
+        GPU_AWARE_MPI="${{ inputs.gpu-aware-mpi }}"
+        if [ "${GPU_AWARE_MPI}" = "true" ]; then
+          source "${GITHUB_WORKSPACE}/.github/scripts/gpu-aware-mpi-env.sh"
+        fi
 
         EXCLUDE_FILE=""
         if [ -n "${ECT_EXCLUDED_VARS}" ]; then
@@ -181,6 +191,10 @@ runs:
           sed -i '/<immutable_stream name="restart"/,/\/>/ s/output_interval="[^"]*"/output_interval="none"/' streams.atmosphere
           sed -i '/<stream name="output"/,/>/ s/output_interval="[^"]*"/output_interval="'"${RUN_DURATION}"'"/' streams.atmosphere
 
+          if [ "${GPU_AWARE_MPI}" = "true" ]; then
+            "${GITHUB_WORKSPACE}/.github/scripts/enable-gpu-aware-mpi.sh" namelist.atmosphere
+          fi
+
           echo "[$(date +%H:%M:%S)] Running MPAS-A (${RUN_DURATION}, ${NRANKS} ranks)..."
           set +e
           timeout ${RUN_TIMEOUT}m mpirun -n ${NRANKS} ${MPI_FLAGS} ./atmosphere_model

.github/scripts/enable-gpu-aware-mpi.sh

@@ -0,0 +1,26 @@
+#!/bin/bash
+# Set config_gpu_aware_mpi = .true. in the &development group of a namelist.atmosphere
+# (creating the group if it is missing). MPAS host-stages halo exchanges itself unless this is set,
+# so the MPI library's GPU support (see gpu-aware-mpi-env.sh) has nothing to do without it.
+# usage: enable-gpu-aware-mpi.sh <path/to/namelist.atmosphere>
+set -e
+NML="${1:?usage: enable-gpu-aware-mpi.sh <namelist.atmosphere>}"
+python3 - "${NML}" <<'PY'
+import re, sys
+from pathlib import Path
+p = Path(sys.argv[1])
+text = p.read_text()
+if re.search(r"^\s*&development\b", text, flags=re.MULTILINE | re.IGNORECASE):
+    def edit(m):
+        block = m.group(0)
+        assignment = r"^([^!\n]*\bconfig_gpu_aware_mpi\s*=\s*)\.\w+\."
+        if re.search(assignment, block, flags=re.MULTILINE | re.IGNORECASE):
+            return re.sub(assignment, r"\1.true.", block, flags=re.MULTILINE | re.IGNORECASE)
+        return re.sub(r"(\n\s*/\s*)$", r"\n    config_gpu_aware_mpi = .true.\1", block, count=1)
+    text = re.sub(r"^\s*&development\b[\s\S]*?\n\s*/\s*", edit, text, count=1, flags=re.MULTILINE | re.IGNORECASE)
+else:
+    text = text.rstrip() + "\n&development\n    config_gpu_aware_mpi = .true.\n/\n"
+p.write_text(text)
+PY
+echo "GPU-aware MPI enabled in ${NML}; &development now reads:"
+awk '/^[[:space:]]*&development/,/^[[:space:]]*\//' "${NML}"

.github/scripts/gpu-aware-mpi-env.sh

@@ -0,0 +1,19 @@
+# Source this (do not execute) to set the MPI-library environment for GPU-aware halos (device buffers handed to MPI).
+# Requires MPI_IMPL (mpich|openmpi); MPI_FLAGS is left untouched. Pair with enable-gpu-aware-mpi.sh (namelist side).
+#
+# Verified in the CI GPU containers (almalinux9-nvhpc-{mpich,openmpi}-cuda-26.02) by the check-gpu-aware-mpi workflow:
+#   mpich   4.3.2, ch4:ucx, --with-cuda: handles device buffers. GPU support is on by default; it is set explicitly below
+#           (MPIR_CVAR_ENABLE_GPU is stock MPICH; MPICH_GPU_SUPPORT_ENABLED is Cray MPICH's variable, harmless here).
+#   openmpi 5.0.9, --with-cuda (accelerator:cuda, btl:smcuda; there is NO ucx pml/osc component in the image): CUDA-aware on a
+#           node by default. Forcing OMPI_MCA_pml=ucx / --mca pml ucx aborts MPI_Init ("Framework: pml, Component: ucx not
+#           found"), so nothing is set for OpenMPI.
+# If a container's MPI is not built with CUDA support, the run fails loudly (segfault/abort) instead of silently staging.
+case "${MPI_IMPL}" in
+  mpich)
+    export MPIR_CVAR_ENABLE_GPU=1
+    export MPICH_GPU_SUPPORT_ENABLED=1
+    ;;
+  openmpi)
+    : # CUDA-aware by default; do not force UCX (see above)
+    ;;
+esac

.github/workflows/_test-bfb.yml

@@ -28,6 +28,11 @@ on:
         description: 'MPI implementation for the container image (mpich, openmpi)'
         required: true
         type: string
+      gpu-aware-mpi:
+        description: 'Use GPU-aware MPI halo exchanges on the GPU (OpenACC) runs (config_gpu_aware_mpi + MPI GPU-support env). Needs CUDA-enabled MPI in the container; set false to fall back to host-staged halos.'
+        required: false
+        type: boolean
+        default: true
       variants:
         description: >
           JSON array of at least two variant objects. Required per variant: id (unique slug),
@@ -276,6 +281,7 @@ jobs:
           mpi-impl: ${{ inputs.mpi }}
           working-dir: run-${{ matrix.id }}
           strict-exit-check: 'false'
+          gpu-aware-mpi: ${{ matrix.openacc && inputs.gpu-aware-mpi }}
 
       - name: Upload history output
         uses: actions/upload-artifact@v6

.github/workflows/_test-gpu.yml

@@ -17,6 +17,11 @@ on:
         description: 'MPI implementation (mpich, openmpi)'
         required: true
         type: string
+      gpu-aware-mpi:
+        description: 'Use GPU-aware MPI halo exchanges on the GPU (OpenACC) runs (config_gpu_aware_mpi + MPI GPU-support env). Needs CUDA-enabled MPI in the container; set false to fall back to host-staged halos.'
+        required: false
+        type: boolean
+        default: true
       # Cross-repo inputs disabled until NCAR security review of running
       # external code on self-hosted CIRRUS runners is complete.
       # mpas-repository:
@@ -199,6 +204,7 @@ jobs:
           run-duration: '0_02:36:00'
           run-timeout: '45'
           restart-file: ${{ steps.restart.outputs.available == 'true' && steps.restart.outputs.file || '' }}
+          gpu-aware-mpi: ${{ inputs.gpu-aware-mpi }}
 
       - name: Upload history file
         uses: actions/upload-artifact@v6

.github/workflows/_test-perf.yml

@@ -28,6 +28,11 @@ on:
         description: 'MPI implementation (mpich, openmpi)'
         required: true
         type: string
+      gpu-aware-mpi:
+        description: 'Use GPU-aware MPI halo exchanges on the GPU (OpenACC) runs (config_gpu_aware_mpi + MPI GPU-support env). Needs CUDA-enabled MPI in the container; set false to fall back to host-staged halos.'
+        required: false
+        type: boolean
+        default: true
       resolution:
         description: 'Test case resolution (see RELEASE_TESTDATA_* in ci-config.env)'
         required: false
@@ -197,6 +202,8 @@ jobs:
           # mpi-run-flags input -- does not touch ci-config.env, so no other
           # workflow that calls run-mpas is affected.
           mpi-run-flags: '-bind-to core'
+          # GPU leg only (the CPU build has no device data); CPU leg unchanged.
+          gpu-aware-mpi: ${{ matrix.openacc && inputs.gpu-aware-mpi }}
 
       - name: Parse MPAS timers
         uses: ./.github/actions/parse-mpas-timers

.github/workflows/bfb-decomp-gpu.yml

@@ -5,11 +5,18 @@ name: "BFB: Decomposition GPU (1 vs 4 ranks)"
 
 on:
   workflow_dispatch:
+    inputs:
+      gpu_aware_mpi:
+        description: 'Use GPU-aware MPI halo exchanges on the GPU (OpenACC) runs. Needs CUDA-enabled MPI in the container; set false to fall back to host-staged halos.'
+        required: false
+        type: boolean
+        default: true
 
 jobs:
   test:
     uses: ./.github/workflows/_test-bfb.yml
     with:
+      gpu-aware-mpi: ${{ inputs.gpu_aware_mpi }}
       compiler: nvhpc
       mpi: mpich
       gpu: 'true'

.github/workflows/bfb-io-gpu.yml

@@ -5,11 +5,18 @@ name: "BFB: I/O GPU (SMIOL vs PIO)"
 
 on:
   workflow_dispatch:
+    inputs:
+      gpu_aware_mpi:
+        description: 'Use GPU-aware MPI halo exchanges on the GPU (OpenACC) runs. Needs CUDA-enabled MPI in the container; set false to fall back to host-staged halos.'
+        required: false
+        type: boolean
+        default: true
 
 jobs:
   test:
     uses: ./.github/workflows/_test-bfb.yml
     with:
+      gpu-aware-mpi: ${{ inputs.gpu_aware_mpi }}
       compiler: nvhpc
       mpi: mpich
       gpu: 'true'

.github/workflows/bfb-nvhpc-cpu-vs-gpu.yml

@@ -9,11 +9,18 @@ name: "BFB: NVHPC CPU vs GPU (OpenACC)"
 
 on:
   workflow_dispatch:
+    inputs:
+      gpu_aware_mpi:
+        description: 'Use GPU-aware MPI halo exchanges on the GPU (OpenACC) runs. Needs CUDA-enabled MPI in the container; set false to fall back to host-staged halos.'
+        required: false
+        type: boolean
+        default: true
 
 jobs:
   test:
     uses: ./.github/workflows/_test-bfb.yml
     with:
+      gpu-aware-mpi: ${{ inputs.gpu_aware_mpi }}
       compiler: nvhpc
       mpi: mpich
       gpu: 'false'

.github/workflows/perf-cpu-vs-gpu.yml

@@ -14,6 +14,11 @@ name: "Perf: CPU vs GPU (NVHPC)"
 on:
   workflow_dispatch:
     inputs:
+      gpu_aware_mpi:
+        description: 'Use GPU-aware MPI halo exchanges on the GPU (OpenACC) runs. Needs CUDA-enabled MPI in the container; set false to fall back to host-staged halos.'
+        required: false
+        type: boolean
+        default: true
       mpi:
         description: MPI implementation
         required: true
@@ -62,6 +67,7 @@ jobs:
   perf:
     uses: ./.github/workflows/_test-perf.yml
     with:
+      gpu-aware-mpi: ${{ inputs.gpu_aware_mpi }}
       mpi: ${{ inputs.mpi }}
       precision: ${{ inputs.precision }}
       resolution: ${{ inputs.resolution }}

.github/workflows/test-gpu-mpich.yml

@@ -6,10 +6,17 @@ name: "NVHPC+MPICH (GPU)"
 
 on:
   workflow_dispatch:
+    inputs:
+      gpu_aware_mpi:
+        description: 'Use GPU-aware MPI halo exchanges on the GPU (OpenACC) runs. Needs CUDA-enabled MPI in the container; set false to fall back to host-staged halos.'
+        required: false
+        type: boolean
+        default: true
 
 jobs:
   test:
     uses: ./.github/workflows/_test-gpu.yml
     with:
+      gpu-aware-mpi: ${{ inputs.gpu_aware_mpi }}
       mpi: mpich
 

.github/workflows/test-gpu-openmpi.yml

@@ -7,10 +7,17 @@ name: "NVHPC+OpenMPI (GPU)"
 
 on:
   workflow_dispatch:
+    inputs:
+      gpu_aware_mpi:
+        description: 'Use GPU-aware MPI halo exchanges on the GPU (OpenACC) runs. Needs CUDA-enabled MPI in the container; set false to fall back to host-staged halos.'
+        required: false
+        type: boolean
+        default: true
 
 jobs:
   test:
     uses: ./.github/workflows/_test-gpu.yml
     with:
+      gpu-aware-mpi: ${{ inputs.gpu_aware_mpi }}
       mpi: openmpi