diff --git a/.github/actions/run-mpas/action.yml b/.github/actions/run-mpas/action.yml index 9c9f79cda2..c28cd6d5d6 100644 --- a/.github/actions/run-mpas/action.yml +++ b/.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,6 +162,10 @@ 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 ===" @@ -156,6 +173,7 @@ runs: 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)" diff --git a/.github/actions/run-perturb-mpas/action.yml b/.github/actions/run-perturb-mpas/action.yml index 9af1660708..df5952a1f0 100644 --- a/.github/actions/run-perturb-mpas/action.yml +++ b/.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 '// s/output_interval="[^"]*"/output_interval="none"/' streams.atmosphere sed -i '// 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 diff --git a/.github/scripts/enable-gpu-aware-mpi.sh b/.github/scripts/enable-gpu-aware-mpi.sh new file mode 100755 index 0000000000..a0eb6a6007 --- /dev/null +++ b/.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 +set -e +NML="${1:?usage: enable-gpu-aware-mpi.sh }" +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}" diff --git a/.github/scripts/gpu-aware-mpi-env.sh b/.github/scripts/gpu-aware-mpi-env.sh new file mode 100644 index 0000000000..657553a31e --- /dev/null +++ b/.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 diff --git a/.github/workflows/_test-bfb.yml b/.github/workflows/_test-bfb.yml index aa5ff3732c..ad71ab719a 100644 --- a/.github/workflows/_test-bfb.yml +++ b/.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 diff --git a/.github/workflows/_test-gpu.yml b/.github/workflows/_test-gpu.yml index 4bcd07db07..4c44ffa426 100644 --- a/.github/workflows/_test-gpu.yml +++ b/.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 diff --git a/.github/workflows/_test-perf.yml b/.github/workflows/_test-perf.yml index 54d5149b06..43414f9377 100644 --- a/.github/workflows/_test-perf.yml +++ b/.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 diff --git a/.github/workflows/bfb-decomp-gpu.yml b/.github/workflows/bfb-decomp-gpu.yml index a3bd6e45d7..f5af2364b0 100644 --- a/.github/workflows/bfb-decomp-gpu.yml +++ b/.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' diff --git a/.github/workflows/bfb-io-gpu.yml b/.github/workflows/bfb-io-gpu.yml index 34fe2967cb..a8aef4347f 100644 --- a/.github/workflows/bfb-io-gpu.yml +++ b/.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' diff --git a/.github/workflows/bfb-nvhpc-cpu-vs-gpu.yml b/.github/workflows/bfb-nvhpc-cpu-vs-gpu.yml index 78d477a7e4..6995d248b1 100644 --- a/.github/workflows/bfb-nvhpc-cpu-vs-gpu.yml +++ b/.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' diff --git a/.github/workflows/perf-cpu-vs-gpu.yml b/.github/workflows/perf-cpu-vs-gpu.yml index 2fd23dc1f2..7e240a5676 100644 --- a/.github/workflows/perf-cpu-vs-gpu.yml +++ b/.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 }} diff --git a/.github/workflows/test-gpu-mpich.yml b/.github/workflows/test-gpu-mpich.yml index efd1097fa6..619fc1b7c2 100644 --- a/.github/workflows/test-gpu-mpich.yml +++ b/.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 diff --git a/.github/workflows/test-gpu-openmpi.yml b/.github/workflows/test-gpu-openmpi.yml index fff0259dcf..52ca879cf4 100644 --- a/.github/workflows/test-gpu-openmpi.yml +++ b/.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